feat(ml): add stateless bundle-local size-aware batching and benchmark (#37532)

* feat(ml): add stateless bundle-local size-aware batching and benchmark

* fix(ml): improve test coverage for SortAndBatchElements

- Exclude *_benchmark.py from codecov (standalone scripts, not production code)
- Remove redundant validation from internal DoFn classes (already validated by PTransform)
- Add direct in-process unit tests for DoFn internals to capture coverage
  (FnApiRunner runs DoFns in separate process, invisible to coverage tools)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Address PR review: clarify benchmark comment and warn on len() fallback

Reframe benchmark docstring to clarify that sorting combined with
weight-based splitting drives the improvement. Move default element
size fallback into DoFn instances with a one-time warning when len()
is unsupported, so users know to provide a custom element_size_fn.

* Migrate JupyterLab sidepanel extension to prebuilt package distribution

Replace deprecated jupyter labextension install/link workflow with
pip-installable prebuilt extension for JupyterLab 4+ compatibility.

- Add install.json for prebuilt extension discovery metadata
- Add style/index.js CSS entry point and styleModule field in package.json
- Include js in package.json files glob so style/index.js is published
- Add Extensions and Extensions :: Prebuilt classifiers to pyproject.toml
- Add missing src/yaml/* to tsconfig.json includes
- Remove deprecated labextension install/link/build instructions from READMEs
- Replace ipywidgets labextension install with pip install in Interactive README

* Use real Beam pipelines in sort-and-batch benchmark

* Address all review comments

- Reuse _WindowAwareBatchingDoFn._MAX_LIVE_WINDOWS instead of keeping a separate hard-coded limit in SortAndBatchElements.
- Drop the padding-efficiency unit test that compared incongruent batching strategies and keep the transform tests focused on deterministic behavior.
- Align benchmark typing with modern Python style by using collections.abc imports and native built-in generics.
- Make the sorted-order test clearer by naming the expected batch contents explicitly.

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: Eliaaazzz

.github/codecov.yml

@@ -73,6 +73,7 @@ ignore:
   - "**/*_microbenchmark.py"
   - "sdks/go/pkg/beam/register/register.go"
   - "sdks/python/apache_beam/testing/benchmarks/nexmark/**"
+  - "**/*_benchmark.py"
   - "sdks/python/apache_beam/examples/**"
 
 # See https://docs.codecov.com/docs/flags for options.

sdks/python/apache_beam/runners/interactive/README.md

@@ -244,23 +244,10 @@ a quick reference). For a more general and complete getting started guide, see
     jupyter kernelspec list
     ```
 
-*   Extend JupyterLab through labextension. **Note**: labextension is different from nbextension
-    from pre-lab jupyter notebooks.
-
-    All jupyter labextensions need nodejs
-
-    ```bash
-    # Homebrew users do
-    brew install node
-    # Or Conda users do
-    conda install -c conda-forge nodejs
-    ```
-
-    Enable ipywidgets
+*   Install ipywidgets (includes the JupyterLab widget manager as a prebuilt extension):
 
     ```bash
     pip install ipywidgets
-    jupyter labextension install @jupyter-widgets/jupyterlab-manager
     ```
 
 ### Start the notebook

sdks/python/apache_beam/runners/interactive/extensions/apache-beam-jupyterlab-sidepanel/README.md

@@ -31,41 +31,22 @@ Includes two different side panels:
 
 ## Installation
 
-There are two ways to install the extension:
-
-### 1. Via pip (recommended)
-
-The extension is now available as a Python package on PyPI. You can install it with:
+This extension is distributed as a prebuilt Python package. Install it with pip:
 
 ```bash
 pip install apache-beam-jupyterlab-sidepanel
 ```
 
-After installation, rebuild JupyterLab to activate the extension:
-
-```bash
-jupyter lab clean
-jupyter lab build
-```
-
-Then restart JupyterLab. The side panels will be available automatically.
+Then restart JupyterLab. The side panels will be available automatically — no
+`jupyter lab build` step is needed.
 
-
-### 2. Via JupyterLab Extension Manager (legacy, will be deprecated soon)
+You can verify the extension is installed:
 
 ```bash
-jupyter labextension install apache-beam-jupyterlab-sidepanel
+jupyter labextension list
 ```
 
-This installs the extension using JupyterLab's legacy extension system.
-
----
-
-## Notes
-
-- Pip installation is now the preferred method as it handles Python packaging and JupyterLab extension registration seamlessly.
-- After any upgrade or reinstallation, always rebuild JupyterLab to ensure the extension is activated.
-- For detailed usage and development, refer to the source code and issues on [GitHub](https://github.com/apache/beam).
+The extension should appear under the **prebuilt extensions** section.
 
 ---
 
@@ -90,15 +71,12 @@ The `jlpm` command is JupyterLab's pinned version of
 
 # Install dependencies
 jlpm
-# Build Typescript source
-jlpm build
-# Link your development version of the extension with JupyterLab
-jupyter labextension link .
 
-# Rebuild Typescript source after making changes
-jlpm build
-# Rebuild JupyterLab after making any changes
-jupyter lab build
+# Install the extension in editable mode (runs an initial JS build)
+pip install -e .
+
+# Verify installation
+jupyter labextension list
 ```
 
 You can watch the source directory and run JupyterLab in watch mode to watch for changes in the extension's source and automatically rebuild the extension and application.
@@ -110,7 +88,7 @@ jlpm watch
 jupyter lab --watch
 ```
 
-Now every change will be built locally and bundled into JupyterLab. Be sure to refresh your browser page after saving file changes to reload the extension (note: you'll need to wait for webpack to finish, which can take 10s+ at times).
+Now every change will be built locally and bundled into JupyterLab. Be sure to refresh your browser page after saving file changes to reload the extension (note: you'll need to wait for the build to finish, which can take 10s+ at times).
 
 ### Test
 
@@ -214,9 +192,5 @@ $PREFIX/share/jupyter/labextensions/apache-beam-jupyterlab-sidepanel/
 ### Uninstall
 
 ```bash
-jupyter labextension uninstall apache-beam-jupyterlab-sidepanel
-```
-or
-```bash
-pip uninstall apache-beam-jupyterlab-sidepanel
+pip uninstall apache_beam_jupyterlab_sidepanel
 ```

sdks/python/apache_beam/runners/interactive/extensions/apache-beam-jupyterlab-sidepanel/install.json

@@ -0,0 +1,5 @@
+{
+  "packageManager": "python",
+  "packageName": "apache_beam_jupyterlab_sidepanel",
+  "uninstallInstructions": "Use your Python package manager (pip, conda, etc.) to uninstall the package apache_beam_jupyterlab_sidepanel"
+}

sdks/python/apache_beam/runners/interactive/extensions/apache-beam-jupyterlab-sidepanel/package.json

@@ -15,7 +15,7 @@
   "author": "apache-beam",
   "files": [
     "lib/**/*.{d.ts,eot,gif,html,jpg,js,js.map,json,png,svg,woff2,ttf}",
-    "style/**/*.{css,eot,gif,html,jpg,json,png,svg,woff2,ttf}"
+    "style/**/*.{css,js,eot,gif,html,jpg,json,png,svg,woff2,ttf}"
   ],
   "main": "lib/index.js",
   "types": "lib/index.d.ts",
@@ -100,6 +100,7 @@
     "style/*.css",
     "style/index.js"
   ],
+  "styleModule": "style/index.js",
   "jupyterlab": {
     "extension": true,
     "outputDir": "apache_beam_jupyterlab_sidepanel/labextension"

sdks/python/apache_beam/runners/interactive/extensions/apache-beam-jupyterlab-sidepanel/pyproject.toml

@@ -33,6 +33,8 @@ classifiers = [
     "Framework :: Jupyter",
     "Framework :: Jupyter :: JupyterLab",
     "Framework :: Jupyter :: JupyterLab :: 4",
+    "Framework :: Jupyter :: JupyterLab :: Extensions",
+    "Framework :: Jupyter :: JupyterLab :: Extensions :: Prebuilt",
     "License :: OSI Approved :: Apache Software License",
     "Programming Language :: Python :: 3",
 ]

sdks/python/apache_beam/runners/interactive/extensions/apache-beam-jupyterlab-sidepanel/style/index.js

@@ -0,0 +1,13 @@
+// Licensed under the Apache License, Version 2.0 (the 'License'); you may not
+// use this file except in compliance with the License. You may obtain a copy of
+// the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an 'AS IS' BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+
+import './index.css';

sdks/python/apache_beam/runners/interactive/extensions/apache-beam-jupyterlab-sidepanel/tsconfig.json

@@ -29,6 +29,7 @@
     "src/common/*",
     "src/kernel/*",
     "src/inspector/*",
+    "src/yaml/*",
     "src/__tests__/**/*"
   ]
 }