diff --git a/README.md b/README.md index 877ea88f..d5f99222 100644 --- a/README.md +++ b/README.md @@ -53,7 +53,11 @@ pip install weightslab ``` **2. Launch Studio** ```bash +<<<<<<< HEAD +weightslab launch # then open https://localhost:5173 πŸš€ +======= weightslab start # preferred localhost:50051; falls back to a free port if busy πŸš€ +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b ``` **3. Wrap your training script** ```python @@ -116,7 +120,11 @@ For a detailed installation guide and advanced configuration → [Installatio 3. **Run your script, then launch the UI in a separate terminal:** ```bash python train.py +<<<<<<< HEAD + weightslab launch +======= weightslab start +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b ``` 4. **Open your browser** at the URL printed by `weightslab start` and inspect your training in real time. diff --git a/docs/_static/examples-gallery.js b/docs/_static/examples-gallery.js index c5a6a609..d610919b 100644 --- a/docs/_static/examples-gallery.js +++ b/docs/_static/examples-gallery.js @@ -20,7 +20,7 @@ desc: 'CNN digit classifier on MNIST. Register hyperparameters, monitor per-sample loss, and use the deny-aware sampler to focus on hard examples.', tags: ['classification', 'supervised', 'mnist', 'cnn'], url: 'examples/pytorch/classification.html', - colab: COLAB + 'PyTorch/ws-classification.ipynb' + colab: COLAB + 'PyTorch/wl-classification.ipynb' }, { badge: 'PyTorch', color: 'pytorch', @@ -28,7 +28,7 @@ desc: 'Per-pixel semantic segmentation with a UNet. Track per-sample IoU and visualise mask overlays directly in the studio.', tags: ['segmentation', 'semantic', 'bdd100k', 'masks', 'dense prediction'], url: 'examples/pytorch/segmentation.html', - colab: COLAB + 'PyTorch/ws-segmentation.ipynb' + colab: COLAB + 'PyTorch/wl-segmentation.ipynb' }, { badge: 'PyTorch', color: 'pytorch', @@ -36,7 +36,7 @@ desc: 'Bounding-box detection on Penn-Fudan pedestrians. Per-instance multi-index dataframe with (sample_id, annotation_id) keys.', tags: ['detection', 'object detection', 'bounding boxes', 'penn-fudan'], url: 'examples/pytorch/detection.html', - colab: COLAB + 'PyTorch/ws-detection.ipynb' + colab: COLAB + 'PyTorch/wl-detection.ipynb' }, { badge: 'PyTorch', color: 'pytorch', @@ -44,7 +44,7 @@ desc: 'Metric learning with triplet loss on face datasets. Store and explore high-dimensional embeddings per sample in the studio.', tags: ['clustering', 'unsupervised', 'embeddings', 'face recognition', 'metric learning'], url: 'examples/pytorch/clustering.html', - colab: COLAB + 'PyTorch/ws-clustering.ipynb' + colab: COLAB + 'PyTorch/wl-clustering.ipynb' }, { badge: 'PyTorch', color: 'pytorch', @@ -52,7 +52,7 @@ desc: 'Unsupervised anomaly detection on MVTec with a multi-task UNet. Monitor reconstruction quality and per-sample anomaly scores.', tags: ['anomaly detection', 'generation', 'unsupervised', 'mvtec', 'reconstruction'], url: 'examples/pytorch/generation.html', - colab: COLAB + 'PyTorch/ws-generation.ipynb' + colab: COLAB + 'PyTorch/wl-generation.ipynb' }, { badge: 'Lightning', color: 'lightning', @@ -81,7 +81,7 @@ desc: 'Dynamic subscribed signal that classifies each sample\'s loss trajectory (monotonic, U-shape, spiked, …) and auto-tags it.', tags: ['loss analysis', 'signal', 'categorical tag', 'per-sample', 'trajectory'], url: 'examples/usecases/loss_shape_classification.html', - colab: COLAB + 'Usecases/ws-segmentation-loss-shapes.ipynb' + colab: COLAB + 'Usecases/wl-segmentation-loss-shapes-classification.ipynb' } ]; diff --git a/docs/_static/wl-ribbon.js b/docs/_static/wl-ribbon.js index 271516b4..cbca2a4f 100644 --- a/docs/_static/wl-ribbon.js +++ b/docs/_static/wl-ribbon.js @@ -10,7 +10,11 @@ 'The studio streams signals in real-time β€” no need to wait for an epoch to end to see results.', 'weightslab start example --cls launches a full MNIST classification demo in one command.', 'Use subscribe_to= on a signal to build reactive per-sample analytics derived from other signals.', +<<<<<<< HEAD + 'Run weightslab launch --certs to enable HTTPS + mTLS for secure remote studio access.', +======= 'Run weightslab start --certs to enable HTTPS + mTLS for secure remote studio access.', +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b 'Set preload_labels=False for large datasets to speed up startup; labels are loaded lazily.', 'Use array_return_proxies=True (default) to avoid loading the full dataset array into RAM.', 'Set WEIGHTSLAB_LOG_LEVEL=DEBUG to see full gRPC logs when debugging connectivity issues.', diff --git a/docs/configuration.rst b/docs/configuration.rst index d699e536..17f5c201 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -354,7 +354,11 @@ Deploying the Studio ~~~~~~~~~~~~~~~~~~~~~ All Weights Studio configuration variables are passed to the UI at launch time +<<<<<<< HEAD +via ``weightslab launch``. There are two ways to supply them. +======= via ``weightslab start``. There are two ways to supply them. +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b **Option 1 β€” shell exports (quick, per-session)** @@ -362,7 +366,11 @@ via ``weightslab start``. There are two ways to supply them. export ENABLE_AGENT=0 export BB_THUMB_RENDER=50 +<<<<<<< HEAD + weightslab launch +======= weightslab start +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b **Option 2 β€” ``.env`` file (persistent, version-controllable)** @@ -380,7 +388,11 @@ Then launch normally: .. code-block:: bash +<<<<<<< HEAD + weightslab launch +======= weightslab start +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b WeightsLab loads the ``.env`` file automatically. Shell exports take precedence over ``.env`` values. diff --git a/docs/examples/pytorch/classification.rst b/docs/examples/pytorch/classification.rst index c5a3cfad..229351ce 100644 --- a/docs/examples/pytorch/classification.rst +++ b/docs/examples/pytorch/classification.rst @@ -157,14 +157,18 @@ debug values, custom distances, etc. .. code-block:: bash +<<<<<<< HEAD + weightslab launch # 1. deploy the studio +======= weightslab start # 1. deploy the studio +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b weightslab start example --cls # 2. start the classification demo .. raw:: html
- + Open In Colab
diff --git a/docs/examples/pytorch/clustering.rst b/docs/examples/pytorch/clustering.rst index cb4d9eb1..c6725940 100644 --- a/docs/examples/pytorch/clustering.rst +++ b/docs/examples/pytorch/clustering.rst @@ -93,14 +93,18 @@ silhouette score to identify consistently confused identities. .. code-block:: bash +<<<<<<< HEAD + weightslab launch # 1. deploy the studio +======= weightslab start # 1. deploy the studio +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b weightslab start example --clus # 2. start the clustering demo .. raw:: html
- + Open In Colab
diff --git a/docs/examples/pytorch/detection.rst b/docs/examples/pytorch/detection.rst index 31873896..a2a37f6b 100644 --- a/docs/examples/pytorch/detection.rst +++ b/docs/examples/pytorch/detection.rst @@ -125,14 +125,18 @@ pipeline. See :ref:`good-practice-get-items` for the recommended signature. .. code-block:: bash +<<<<<<< HEAD + weightslab launch # 1. deploy the studio +======= weightslab start # 1. deploy the studio +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b weightslab start example --det # 2. start the detection demo .. raw:: html
- + Open In Colab
diff --git a/docs/examples/pytorch/generation.rst b/docs/examples/pytorch/generation.rst index 12a99013..4750317b 100644 --- a/docs/examples/pytorch/generation.rst +++ b/docs/examples/pytorch/generation.rst @@ -99,5 +99,9 @@ filter by pair distance to find the hardest negatives. .. code-block:: bash +<<<<<<< HEAD + weightslab launch # 1. deploy the studio +======= weightslab start # 1. deploy the studio +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b weightslab start example --gen # 2. start the generation / anomaly demo diff --git a/docs/examples/pytorch/segmentation.rst b/docs/examples/pytorch/segmentation.rst index caac785b..ab4a2519 100644 --- a/docs/examples/pytorch/segmentation.rst +++ b/docs/examples/pytorch/segmentation.rst @@ -112,14 +112,18 @@ these signals over training steps and lets you sort samples by them. .. code-block:: bash +<<<<<<< HEAD + weightslab launch # 1. deploy the studio +======= weightslab start # 1. deploy the studio +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b weightslab start example --seg # 2. start the segmentation demo .. raw:: html
- + Open In Colab
diff --git a/docs/examples/usecases/lidar_detection.rst b/docs/examples/usecases/lidar_detection.rst index b86b13a5..ef32aead 100644 --- a/docs/examples/usecases/lidar_detection.rst +++ b/docs/examples/usecases/lidar_detection.rst @@ -120,7 +120,11 @@ and override ``load_points`` and optionally ``render_thumbnail_2d``: .. code-block:: bash +<<<<<<< HEAD + weightslab launch # 1. deploy the studio +======= weightslab start # 1. deploy the studio +>>>>>>> a6f0c1acbac49221358a21b0bde87b348c09204b weightslab start example --2d_det # 2a. 2D pillar-grid detection # or weightslab start example --3d_det # 2b. 3D bounding-box detection diff --git a/docs/examples/usecases/loss_shape_classification.rst b/docs/examples/usecases/loss_shape_classification.rst index d7ad717e..d5240632 100644 --- a/docs/examples/usecases/loss_shape_classification.rst +++ b/docs/examples/usecases/loss_shape_classification.rst @@ -213,7 +213,7 @@ Workflow in the studio .. raw:: html
- + Open In Colab
diff --git a/docs/four_way_approach.rst b/docs/four_way_approach.rst index 9fda43b8..effa9d2d 100644 --- a/docs/four_way_approach.rst +++ b/docs/four_way_approach.rst @@ -53,6 +53,6 @@ Typical integration flow to custom Python script .. code-block:: bash weightslab start - python train.py + python train.py - Resume training from the UI and use tags/discards/signals to iteratively improve data and model behavior. diff --git a/docs/quickstart.rst b/docs/quickstart.rst index e3e9eefe..d0ee011d 100644 --- a/docs/quickstart.rst +++ b/docs/quickstart.rst @@ -139,7 +139,7 @@ Pass ``--certs`` to generate (if missing) and use TLS certificates + a gRPC auth .. code-block:: bash weightslab start # unsecured HTTP (default) - weightslab start --certs # secured HTTPS + gRPC auth (run `weightslab se` first) + weightslab start --certs # secured HTTPS + gRPC auth (run `weightslab se` first) .. important:: diff --git a/docs/usage/docker.rst b/docs/usage/docker.rst new file mode 100644 index 00000000..c32ae292 --- /dev/null +++ b/docs/usage/docker.rst @@ -0,0 +1,350 @@ +.. _docker-usage: + +Docker Usage +=============== + +**Examples:** ``weightslab/examples/Docker_training/`` + +WeightsLab supports running the training script β€” and the entire UI stack β€” +inside Docker. Three integration patterns are available; they differ in **how +the trainer container gets a Docker daemon** to launch the Envoy proxy and the +Weights Studio frontend alongside it. + +.. list-table:: + :header-rows: 1 + :widths: 28 24 24 24 + + * - + - **A Β· Docker-in-Docker** + - **B Β· Siblings (DooD)** + - **C Β· Self-contained siblings** + * - Docker daemon + - Own nested daemon + - Host daemon (socket mount) + - Host daemon (socket mount) + * - Envoy/frontend run + - Nested inside trainer + - Siblings on the host + - Siblings on the host + * - Starts UI via + - ``weightslab launch`` + - ``weightslab launch`` + - Custom ``ui-compose.yml`` + * - ``--privileged`` + - **Required** + - No + - No + * - gRPC ``:50051`` published + - Not published + - Yes + - Yes + * - Host bind mounts + - None (co-located fs) + - Path alignment required + - **None** + * - Host setup + - None + - ``setup-host.sh`` + - **None** + * - HTTPS (optional) + - ``WEIGHTSLAB_TLS=1`` + - ``WEIGHTSLAB_TLS=1`` + - ``WEIGHTSLAB_TLS=1`` + * - Windows / Docker Desktop + - Yes + - Awkward (use WSL2) + - **Yes** + +.. note:: + + - **A (DinD)** β€” fully self-contained; needs ``--privileged``. Best when + you want isolation or are on Windows. + - **B (DooD)** β€” uses the stock ``weightslab launch``; cleanest on a + Linux host with path alignment already done. + - **C (self-contained siblings)** β€” no host prep, no bind mounts, + works natively on Windows. + +.. _docker-dind: + +Option A: Docker-in-Docker (DinD) +---------------------------------- + +**Source:** ``weightslab/examples/Docker_training/docker_in_docker/`` + +The trainer container starts its **own inner Docker daemon** and uses +``weightslab launch`` inside it. Because the inner daemon shares the +trainer container's filesystem, all paths resolve without any host-side setup. + +Wiring diagram +~~~~~~~~~~~~~~ + +.. code-block:: text + + host browser + β†’ localhost:5173 ── (re-published) ──► trainer ──► inner frontend :5173 + β†’ localhost:8080 ── (re-published) ──► trainer ──► inner Envoy :8080 + β”‚ grpc-backend:host-gateway + β–Ό + in-process gRPC backend :50051 + (same container as trainer) + +Configuration requirements +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 5 40 30 25 + + * - # + - Requirement + - Location + - Why + * - 1 + - ``privileged: true`` on the trainer service + - ``docker-compose.yml`` + - inner ``dockerd`` cannot run otherwise + * - 2 + - Start the inner ``dockerd`` + wait for socket + - ``entrypoint.sh`` + - nested daemon that hosts Envoy + frontend + * - 3 + - Persist ``/var/lib/docker`` with a named volume + - ``docker-compose.yml`` + - caches image pulls across runs + * - 4 + - Publish ``5173`` + ``8080`` to the host + - ``docker-compose.yml`` + - browser reaches the *nested* containers + * - 5 + - ``GRPC_BACKEND_PORT=50051`` + - ``docker-compose.yml`` + - Envoy's ``grpc-backend:host-gateway`` dials this port + * - 6 + - ``WEIGHTSLAB_SKIP_DOCKER_OPS=1`` before ``ui launch`` + - ``entrypoint.sh`` + - skips the in-container rebuild; pulls the image only + * - 7 + - Order: ``ui launch`` β†’ ``start example`` + - ``entrypoint.sh`` + - UI stack must be up before the backend starts + +GPU access +~~~~~~~~~~ + +Add the ``deploy`` block to the trainer service in ``docker-compose.yml``: + +.. code-block:: yaml + + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + +The host also needs the **NVIDIA Container Toolkit** installed (``sudo +nvidia-ctk runtime configure --runtime=docker``). Comment out the ``deploy`` +block on hosts without an NVIDIA GPU/toolkit β€” ``docker compose up`` will +otherwise fail with *"could not select device driver nvidia"*. + +Running +~~~~~~~ + +.. code-block:: bash + + # from weightslab/examples/Docker_training/docker_in_docker/ + docker compose up --build + +Open http://localhost:5173. Stop with ``Ctrl+C`` or ``docker compose down``. + +.. dropdown:: Enable HTTPS / mTLS (optional) + :color: secondary + + Only needed for remote or production access. For local development, plain + HTTP at http://localhost:5173 works without any certificates. + + DinD is the simplest TLS option because certs, Envoy, and the backend are + co-located: + + .. code-block:: bash + + WEIGHTSLAB_TLS=1 docker compose up --build + + Then trust the generated CA on the host browser (once): + + .. code-block:: powershell + + # Windows β€” pull the CA out of the running container and trust it + docker cp weightslab_trainer_dind:/root/.weightslab-certs/ca.crt . + Import-Certificate -FilePath .\ca.crt -CertStoreLocation Cert:\CurrentUser\Root + + Open https://localhost:5173. + + What ``WEIGHTSLAB_TLS=1`` does in ``entrypoint.sh``: + + 1. Runs ``weightslab launch --certs``, generating certs into + ``WEIGHTSLAB_CERTS_DIR`` and configuring Envoy + the frontend for HTTPS. + 2. Exports ``GRPC_TLS_ENABLED=1`` + ``GRPC_TLS_CERT_DIR`` so the gRPC + backend also speaks TLS (required to prevent Envoy upstream 503s). + +.. _docker-siblings-c: + +Option C: Self-contained siblings (recommended for Windows) +------------------------------------------------------------ + +**Source:** ``weightslab/examples/Docker_training/siblings_self_contained/`` + +The trainer mounts the **host Docker socket** and starts Envoy + the frontend +as sibling containers β€” but unlike Option B, it **never bind-mounts a host +path**. The Envoy config (and, for TLS, the certs) are delivered via **named +volumes over the socket**, so this option works on **Windows / Docker Desktop** +with zero host preparation. + +Wiring diagram +~~~~~~~~~~~~~~ + +.. code-block:: text + + host browser + β†’ localhost:5173 ──► frontend container (sibling, HTTP) + β†’ localhost:8080 ──► Envoy container (sibling, config from named volume) + β”‚ grpc-backend:host-gateway β†’ host:50051 + β–Ό + host:50051 ──► trainer's gRPC backend :50051 + +Configuration requirements +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 5 40 30 25 + + * - # + - Requirement + - Location + - Why + * - 1 + - Mount ``/var/run/docker.sock`` + - ``docker-compose.yml`` + - drive the host daemon (start sibling containers) + * - 2 + - Publish ``50051:50051`` from the trainer + - ``docker-compose.yml`` + - sibling Envoy dials it via ``host-gateway`` + * - 3 + - Stage Envoy config into ``wl_envoy_cfg`` named volume + - ``entrypoint.sh`` + - replaces the host bind mount β€” no host path needed + * - 4 + - ``ui-compose.yml`` with no host bind mounts + - ``ui-compose.yml`` + - nothing for the host daemon to resolve on disk + * - 5 + - ``extra_hosts: grpc-backend:host-gateway`` on the Envoy service + - ``ui-compose.yml`` + - routes Envoy β†’ host:50051 β†’ trainer + * - 6 + - ``GRPC_BACKEND_PORT=50051`` + - ``docker-compose.yml`` + - port the backend binds + Envoy dials + +How config delivery works without bind mounts +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Option C does not call ``weightslab launch``. Instead, ``entrypoint.sh``: + +1. Renders Envoy's plaintext config from the installed ``weightslab`` package. +2. Pipes it into the ``wl_envoy_cfg`` named volume via the socket using a + temporary ``busybox`` container β€” no host path involved. +3. Brings up Envoy + frontend using ``ui-compose.yml``, which reads the config + from the named volume rather than a bind-mounted host directory. + +GPU access +~~~~~~~~~~ + +Same as Option A β€” add the ``deploy`` block to the trainer service and install +the NVIDIA Container Toolkit on the host. + +Running +~~~~~~~ + +.. code-block:: bash + + # from weightslab/examples/Docker_training/siblings_self_contained/ + docker compose up --build + +Open http://localhost:5173. + +**Stopping:** + +.. code-block:: bash + + docker compose down # stop trainer + docker compose -p weightslab_ui -f ui-compose.yml down # stop UI siblings + docker volume rm wl_envoy_cfg # optional: drop staged config + +.. dropdown:: Enable HTTPS / mTLS (optional) + :color: secondary + + Only needed for remote or production access. For local development, plain + HTTP at http://localhost:5173 works without any certificates. + + TLS works here without host bind mounts β€” certs are piped into named volumes + the same way as ``envoy.yaml``: + + .. code-block:: bash + + WEIGHTSLAB_TLS=1 docker compose up --build + + Then trust the CA on the host browser (once): + + .. code-block:: powershell + + docker cp weightslab_trainer_selfcontained:/root/.weightslab-certs/ca.crt . + Import-Certificate -FilePath .\ca.crt -CertStoreLocation Cert:\CurrentUser\Root + + Open https://localhost:5173. + + What ``WEIGHTSLAB_TLS=1`` delivers (all via named volumes, no host paths): + + .. list-table:: + :header-rows: 1 + :widths: 35 30 35 + + * - Layer + - What's needed + - How it arrives + * - Browser ↔ Envoy + - ``envoy-server.crt/key`` + - ``wl_envoy_cfg`` volume + * - Envoy ↔ backend (mTLS) + - ``envoy-client.crt/key`` + ``ca.crt`` + - ``wl_envoy_cfg`` volume + * - Browser ↔ frontend + - ``envoy-server.crt/key`` + - ``wl_nginx_certs`` volume + * - Backend gRPC server + - ``backend-server.crt/key`` + ``ca.crt`` + - ``GRPC_TLS_ENABLED=1`` + ``GRPC_TLS_CERT_DIR`` + + .. note:: + + Generated keys are ``0600`` (owned by root). ``entrypoint.sh`` runs + ``chmod a+rX`` on cert files in the volumes so Envoy (non-root) can + read them. + +Common notes +------------ + +- The example starts **paused** (``is_training: false``); start and steer + training from the UI at http://localhost:5173. +- First run is slow: it pulls ``envoyproxy/envoy`` + ``graybx/weightslab`` + and downloads MNIST for the classification example. +- For either option, to build against the dev branch instead of PyPI: + + .. code-block:: bash + + docker compose build \ + --build-arg WEIGHTSLAB_SPEC="git+https://github.com/GrayboxTech/weightslab.git@dev" + docker compose up diff --git a/docs/user_commands.rst b/docs/user_commands.rst index d6fb3f4f..19fd4ab6 100644 --- a/docs/user_commands.rst +++ b/docs/user_commands.rst @@ -29,6 +29,8 @@ Run weightslab, weightslab -h, or weightslab help to print the full built-in hel - Connect to a running experiment interactive console. * - weightslab tunnel - Forward a remote gRPC backend to a local TCP port. + * - weightslab help + - Show the help/banner (same as no command, or -h). weightslab se ~~~~~~~~~~~~~ @@ -37,8 +39,11 @@ weightslab se weightslab se [certs_dir] [--force-certs] -Generates TLS certificates and a gRPC auth token. The cert directory is the -single source of truth for secure mode. +Generates TLS certificates and a gRPC auth token into a certs directory, then +tells you to export ``WEIGHTSLAB_CERTS_DIR`` β€” the **single source of +truth** the training backend, ``weightslab start --certs``, and any new +shell all read to decide whether TLS/auth is on (derived purely from whether +cert files exist in that directory). weightslab start ~~~~~~~~~~~~~~~~ @@ -68,7 +73,6 @@ Examples: weightslab start weightslab start --port 9000 - weightslab start --config ./config.yaml weightslab start --backend-port 50052 weightslab start --certs @@ -79,7 +83,49 @@ weightslab start example weightslab start example [--cls|--seg|--det|--clus|--gen|--3d_det|--2d_det] -Runs a bundled example in the foreground. +Runs one of the bundled PyTorch examples in the foreground (stop with +Ctrl+C). Installs the example's own ``requirements.txt``/``requirements.in`` +first, without prompting, then runs its ``main.py``. + +``weightslab example start [flags]`` (subcommand order swapped) and the bare +``weightslab example`` are accepted as tolerant aliases with identical +behavior β€” they don't appear in ``--help`` on purpose, ``start example`` is +the documented form. + +**Arguments** β€” mutually exclusive; default is ``--cls``: + +.. list-table:: + :header-rows: 1 + + * - Flag + - Example + * - ``--cls`` *(default)* + - Classification + * - ``--seg`` + - Segmentation + * - ``--det`` + - Detection + * - ``--clus`` + - Clustering + * - ``--gen`` + - Generation + * - ``--3d_det`` + - 3D LiDAR point-cloud detection + * - ``--2d_det`` + - 2D LiDAR point-cloud detection + +**Examples** + +.. code-block:: bash + + weightslab start example # classification (default) + weightslab start example --seg # segmentation + weightslab start example --3d_det # 3D LiDAR detection + weightslab example start --det # tolerant alias, same as `start example --det` + +Then, in another terminal: ``weightslab launch`` and open +``http://localhost:5173``. See :doc:`examples/index` for what each example +demonstrates. weightslab cli ~~~~~~~~~~~~~~ @@ -102,9 +148,288 @@ Forwards a remote gRPC endpoint to a local TCP port. - Use raw TCP tunnels (for example bore or ngrok tcp). - Default local listen port is 50051. -Examples: +- The remote tunnel must be **raw TCP**, *not* an HTTP/gRPC-Web tunnel. A + zero-signup option is `bore `_ with its free + public relay: ``bore local 50051 --to bore.pub`` (prints ``bore.pub:``). + ``ngrok tcp 50051`` also works but now requires a credit card on the free tier. +- The backend must run **plaintext** β€” the default ``weightslab launch`` + (no ``--certs``) β€” so no TLS terminates mid-path. + +**Arguments** + +- ``ENDPOINT`` *(positional, optional)* β€” the remote backend as ``host:port`` + (e.g. ``0.tcp.ngrok.io:12345``); a ``tcp://`` prefix is accepted and + stripped. Default: the ``WEIGHTSLAB_TUNNEL_ENDPOINT`` environment variable, so + a bare ``weightslab tunnel`` works once that is exported. +- ``--listen-port``, ``-p`` *(int)* β€” local port to expose. Default: **50051** + (the port the bundled Envoy upstream dials β€” leave it unless you changed + ``GRPC_BACKEND_PORT``). +- ``--listen-host`` *(str)* β€” interface to bind. Default: **auto** β€” + ``127.0.0.1`` on Windows/macOS (Docker Desktop reaches host loopback via + ``host.docker.internal``), ``0.0.0.0`` on Linux (compose ``host-gateway`` + resolves to the bridge IP, which cannot reach a loopback-only listener). +- ``--remote-port`` *(int)* β€” the remote port, when ``ENDPOINT`` has only a + host and no ``:port``. + +**Examples** + +.. code-block:: bash + + weightslab tunnel bore.pub:12345 # bridge remote backend -> localhost:50051 + weightslab tunnel tcp://bore.pub:12345 # tcp:// prefix is fine + weightslab tunnel # uses $WEIGHTSLAB_TUNNEL_ENDPOINT + weightslab tunnel host.example.com --remote-port 50051 + weightslab tunnel host:50051 -p 50055 # expose locally on a different port + +**Typical workflow** (Colab backend, local UI): .. code-block:: bash - weightslab tunnel bore.pub:12345 - weightslab tunnel --listen-port 50052 bore.pub:12345 + # 1) In Colab: expose the training backend over raw TCP (prints bore.pub:) + # !bore local 50051 --to bore.pub + + # 2) On your machine, in two terminals: + weightslab launch # plaintext HTTP (default) + weightslab tunnel bore.pub:12345 # the host:port bore printed + + # 3) Open http://localhost:5173 β€” Studio streams live from Colab. + +.. note:: + + Step 1 can be done for you: call ``wl.serve(serving_grpc=True, + serving_bore=True)`` in the training script. It downloads ``bore``, opens the + relay, and prints the exact ``weightslab tunnel bore.pub:`` line to run + on your machine β€” see ``serve`` in :doc:`user_functions`. + +The command probes the remote on startup (warning, not fatal, if it isn't up +yet), re-resolves the endpoint per connection (so a changing tunnel IP is picked +up), and runs until ``Ctrl+C``. See the classification Colab notebook +(``examples/Notebooks/PyTorch/wl-classification.ipynb``) for the end-to-end +setup. + +.. _cli-console: + +Interactive CLI console +------------------------ + +The console is a local developer REPL for inspecting and controlling a +running experiment through the global ledger. + +- **Transport**: local TCP, plain-text commands, JSON responses. +- **Intended scope**: development / debugging, not a production control plane. +- **Security model**: binds to localhost by default; plain-text protocol + (keep the port private β€” localhost or a private subnet only). + +How to start it +~~~~~~~~~~~~~~~~ + +From your training script (recommended) β€” starts the server; a client REPL +window opens automatically: + +.. code-block:: python + + import weightslab as wl + + wl.serve(serving_grpc=True, serving_cli=True) # serving_cli defaults to True + wl.keep_serving() + +Then, from any other terminal, attach with: + +.. code-block:: bash + + weightslab cli + +To start the server **headless** (no REPL window pops up; attach later on +demand), pass ``spawn_cli_client=False`` β€” see the ``serve`` entry in +:doc:`user_functions`: + +.. code-block:: python + + wl.serve(serving_cli=True, spawn_cli_client=False) + +Low-level equivalents (rarely needed directly β€” ``wl.serve``/``weightslab +cli`` cover the normal workflow): + +.. code-block:: bash + + python -m weightslab.backend.cli serve --host localhost --port 60000 + python -m weightslab.backend.cli client --host localhost --port 60000 + +If no port is given (or port is ``0``), the server picks a free port and +advertises it for auto-discovery. + +Console commands +~~~~~~~~~~~~~~~~~ + +Type ``help`` (or ``h`` / ``?``) inside the console at any time for this same +reference with extra examples. + +Discovery and help +^^^^^^^^^^^^^^^^^^^ + +- ``help`` / ``h`` / ``?`` β€” show all command syntaxes and examples. +- ``status`` β€” compact snapshot: registered models, dataloaders, optimizers, + hyperparameters, and the current model age. +- ``ledger`` / ``ledgers`` / ``snapshot`` β€” same registry snapshot as + ``status``, without the model-age lookup. +- ``dump`` / ``d`` β€” sanitized dump of dataloaders, optimizers, and + hyperparameters (models are omitted to avoid printing huge weight dumps). +- ``ledger_dump`` / ``dump_ledger`` / ``dump_ledger_all`` β€” like ``dump``, + but **includes models** too. Can be large. + +Training control +^^^^^^^^^^^^^^^^^ + +- ``pause`` / ``p`` β€” pause training and set ``is_training=False``. +- ``resume`` / ``r`` β€” resume training and set ``is_training=True``. + +Registry inspection +^^^^^^^^^^^^^^^^^^^^ + +- ``list_models`` β€” registered model names. +- ``list_optimizers`` β€” registered optimizer names. +- ``list_loaders`` / ``loaders`` / ``list_dataloaders`` β€” registered + dataloader names. +- ``plot_model [model_name]`` (aliases: ``plot_arch``, ``plot``) β€” ASCII tree + of the model's architecture. Omit ``model_name`` to use the default + registered model. + +Sample-level dataset operations +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +**Syntax**: ``list_uids [loader_name] [--discarded] [--limit N]`` +(aliases: ``uids``, ``samples``) + +List sample UIDs (with tags and discard status). Omit ``loader_name`` to +check every registered loader; ``--discarded`` restricts to currently +discarded samples; ``--limit N`` caps the count per loader. + +**Syntax**: ``discard [uid2 ...] [--loader loader_name]`` / +``undiscard [uid2 ...] [--loader loader_name]`` + +Mark one or more samples (by sample/UID) as discarded or restore them. Tries +the dataframe-backed path first (equivalent to :func:`discard_samples`); +without ``--loader``, falls back to every registered loader whose dataset +exposes a discard method. + +**Syntax**: ``add_tag [sample_id2 ...] [--loader loader_name]`` +(alias: ``tag``) + +Add a boolean tag to one or more samples. Same dataframe-first, +all-loaders-fallback behavior as ``discard``. + +**Examples** + +.. code-block:: text + + list_uids + list_uids train_loader --discarded + list_uids --limit 20 + discard sample_001 sample_002 + undiscard sample_001 + add_tag sample_001 difficult sample_002 sample_003 + +Hyperparameter operations +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +- ``hp`` (alias: ``hyperparams``) β€” list registered hyperparameter set names. +- ``hp `` β€” show one set's values. ``hp show `` also works. +- ``set_hp [hp_name] `` (aliases: ``sethp``, ``set-hp``) β€” + update one key path. ``hp_name`` may be omitted only when exactly one + hyperparameter set is registered. ``value`` is parsed as JSON first + (so ``32``, ``0.5``, ``true``, ``"a string"`` all work), falling back to + bool/int/float/string coercion. + +**Examples** + +.. code-block:: text + + hp + hp fashion_mnist + set_hp fashion_mnist data.train_loader.batch_size 32 + set_hp optimizer.lr 0.0005 # hp_name omitted β€” only valid with one hp set + +Evaluation +^^^^^^^^^^ + +- ``evaluate [split_name] [--steps N] [--tags tag1,tag2]`` (aliases: ``eval``, + ``ev``) β€” pause training and trigger a background evaluation pass. Default + split: the first registered dataloader. ``--tags`` restricts evaluation to + samples carrying any of the given tags (and implies not using the full + set); ``--steps`` caps the number of batches evaluated. +- ``eval_status`` (aliases: ``es``, ``evaluation_status``) β€” poll progress + of the current evaluation. +- ``cancel_eval`` (aliases: ``ce``, ``cancel_evaluation``) β€” cancel a running + or pending evaluation. + +**Examples** + +.. code-block:: text + + evaluate # default split, full set + evaluate val_loader + evaluate test_loader --steps 50 + evaluate train_loader --tags difficult,outlier + eval_status + cancel_eval + +See the "Evaluation mode" section of :doc:`user_functions` for how this +integrates with (or without) your own training loop. + +Audit mode +^^^^^^^^^^ + +**Syntax**: ``audit [on|off]`` + +Toggles auditor mode: while on, the optimizer's ``step()`` is skipped (the +training loop keeps running and forward/backward still happen) so you can +inspect gradients/activations without modifying weights. With no argument, +prints the current state. + +**Examples** + +.. code-block:: text + + audit on + audit off + audit # show current state + +AI Agent +^^^^^^^^ + +**Syntax**: ``agent ...`` β€” shortcuts: +``query `` / ``ask `` for ``agent query``. + +Initializes and drives the same natural-language agent used by Weights +Studio (discard/tag/sort/analyze via a prompt) from the console. Full +sub-verb reference, examples, and setup: see :doc:`agent`. + +**Examples** + +.. code-block:: text + + agent status + agent init --api-key sk-or-... --model openai/gpt-4o-mini --timeout 20 + agent models + agent model google/gemini-flash-latest + ask tag train samples with loss > 1.2 as goldset + +Session control +^^^^^^^^^^^^^^^^ + +- ``exit`` / ``quit`` β€” close the client connection (handled server-side; + the server replies then closes the socket). +- ``clear`` / ``cls`` β€” clear the local terminal screen. Handled entirely by + the **client**, not sent to the server. + +Developer notes +~~~~~~~~~~~~~~~~ + +- Prefer the console for quick diagnosis and manual interventions; use + Weights Studio for richer visual workflows. +- Keep the CLI port private (localhost, or a private subnet at most) β€” the + protocol is plain text with no authentication. +- Editing hyperparameters is the only supported mutation path for + architecture-level state; there is currently no console command to + freeze/unfreeze layers or resize a model (that lives in the + :doc:`agent` and Weights Studio surfaces, and in the Python API). diff --git a/tests/backend/test_ui_docker_bridge.py b/tests/backend/test_ui_docker_bridge.py new file mode 100644 index 00000000..4778240e --- /dev/null +++ b/tests/backend/test_ui_docker_bridge.py @@ -0,0 +1,732 @@ +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/weightslab/ui_docker_bridge.py b/weightslab/ui_docker_bridge.py new file mode 100644 index 00000000..11a1e826 --- /dev/null +++ b/weightslab/ui_docker_bridge.py @@ -0,0 +1,1319 @@ +"""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()