"""Dask client helpers for distributed aggregate generation."""
from __future__ import annotations
import logging
import os
import subprocess
import tomllib
from pathlib import Path
logger = logging.getLogger(__name__)
# Project root: two levels up from this file (pyfracval/dask_runner.py → project root)
_PROJECT_ROOT = Path(__file__).parent.parent
def _build_wheel() -> Path:
"""Build a wheel for the local package and return its path."""
logger.info("Building pyfracval wheel with 'uv build'…")
result = subprocess.run(
["uv", "build", "--wheel"],
cwd=str(_PROJECT_ROOT),
capture_output=True,
text=True,
)
if result.returncode != 0:
raise RuntimeError(f"'uv build' failed:\n{result.stdout}\n{result.stderr}")
dist_dir = _PROJECT_ROOT / "dist"
wheels = sorted(dist_dir.glob("*.whl"), key=lambda p: p.stat().st_mtime)
if not wheels:
raise FileNotFoundError(f"No wheel found in {dist_dir} after 'uv build'")
wheel = wheels[-1]
logger.info(f"Built wheel: {wheel.name} ({wheel.stat().st_size // 1024} KB)")
return wheel
def _project_version() -> str:
"""Read project version from pyproject.toml."""
pyproject = _PROJECT_ROOT / "pyproject.toml"
with open(pyproject, "rb") as fh:
data = tomllib.load(fh)
return str(data["project"]["version"])
[docs]
def install_wheel_on_workers(
client,
wheel_path: str | Path,
package_name: str,
expected_version: str,
) -> None:
"""Install an arbitrary wheel on a Dask scheduler and all its workers,
verified via a runtime version fingerprint.
Generic version of the mechanism this module has always used for
installing *pyfracval* itself onto Dask Docker workers that don't have
it preinstalled (e.g. a generic ``ghcr.io/dask/dask`` image) -- reusable
for any package, not just this one. In particular: a *compiled*
extension (e.g. pyfastmm's f2py extension) needs a wheel actually built
for the worker's platform/Python ABI -- this function only ships and
installs whatever wheel you hand it, it does not build one (see
``_register_package`` below for pyfracval's own "build one first" case).
Installs via a ``WorkerPlugin`` (``client.register_plugin()``), not a
one-shot ``client.run()`` sweep: a plugin also runs its ``setup()`` on
any worker that joins *after* this call (Docker restart, autoscaling,
...), where a ``client.run()`` snapshot of currently-connected workers
would silently miss it and leave that worker without the package. The
installer function and plugin class are both defined inline so
cloudpickle serialises them **by value** (bytecode), not by reference to
a module -- which would fail on the scheduler/workers before
*package_name* is installed there.
Parameters
----------
client:
A connected ``dask.distributed.Client``.
wheel_path:
Path to a ``.whl`` file, already built for the *worker's* platform
and Python version (this function does no cross-compilation or
compatibility checking -- get that part right before calling this).
package_name:
The distribution/import name (e.g. ``"pyfracval"``, ``"pyfastmm"``,
``"spcwth"``) -- used for the post-install version check, the
stale-module cache eviction, and the ``<PACKAGE>_INSTALLED_WHEEL``/
``<PACKAGE>_EXPECTED_VERSION`` env vars set on each worker.
expected_version:
Version string the installed wheel must report after installing,
or this raises ``RuntimeError``.
"""
wheel_path = Path(wheel_path)
wheel_bytes = wheel_path.read_bytes()
wheel_filename = wheel_path.name
env_installed_wheel = f"{package_name.upper()}_INSTALLED_WHEEL"
env_expected_version = f"{package_name.upper()}_EXPECTED_VERSION"
logger.info(
f"Installing {wheel_filename} ({len(wheel_bytes) // 1024} KB) "
f"on scheduler and all workers…"
)
def _install_wheel_bytes_embedded(
wheel_bytes: bytes,
wheel_filename: str,
package_name: str,
expected_version: str,
env_installed_wheel: str,
env_expected_version: str,
) -> dict[str, str]:
import importlib
import os
import subprocess
import sys
import tempfile
tmp_dir = tempfile.mkdtemp(prefix=f"{package_name}_wheel_")
wheel_path = os.path.join(tmp_dir, wheel_filename)
with open(wheel_path, "wb") as fh:
fh.write(wheel_bytes)
install_cmds = [
[
"uv",
"pip",
"install",
"--python",
sys.executable,
"--force-reinstall",
wheel_path,
],
[
sys.executable,
"-m",
"pip",
"install",
"--force-reinstall",
wheel_path,
],
]
last_exc: Exception | None = None
for cmd in install_cmds:
try:
subprocess.check_call(cmd)
last_exc = None
break
except Exception as exc: # pragma: no cover - env-specific
last_exc = exc
if last_exc is not None:
raise RuntimeError(
"Failed to install wheel on scheduler/worker using pip and uv pip"
) from last_exc
importlib.invalidate_caches()
for mod_name in list(sys.modules):
if mod_name == package_name or mod_name.startswith(package_name + "."):
sys.modules.pop(mod_name, None)
os.environ[env_installed_wheel] = wheel_filename
os.environ[env_expected_version] = expected_version
return {
"python": sys.executable,
"pid": str(os.getpid()),
"installed_wheel": wheel_filename,
"expected_version": expected_version,
}
def _worker_fingerprint_embedded(package_name: str) -> dict[str, str]:
import importlib
import os
from importlib.metadata import PackageNotFoundError
from importlib.metadata import version as pkg_version
try:
module = importlib.import_module(package_name)
except ImportError:
module = None
try:
runtime_version = pkg_version(package_name)
except PackageNotFoundError:
runtime_version = "unknown"
return {
"version": str(runtime_version),
"module_file": str(getattr(module, "__file__", "unknown")),
"pid": str(os.getpid()),
}
# A WorkerPlugin, not just a one-shot client.run() sweep: client.run()
# only ever reaches workers connected *at call time*. A worker that
# joins later (Docker restart, --nworkers auto scaling up, ...) would
# never get the package and would die with a bare ModuleNotFoundError
# the moment it's handed a task referencing it. register_plugin()
# installs on every currently-connected worker (synchronously -- it
# waits for setup() to finish on each one and returns their results)
# *and* registers with the scheduler so any future worker runs setup()
# automatically on connect. This is the gap client.run() can't close.
#
# Defined nested (like _install_wheel_bytes_embedded above) so
# cloudpickle ships the class by value, not by reference to this
# module -- a fresh worker doesn't have pyfracval importable yet,
# which is exactly the problem this plugin exists to fix.
from distributed import WorkerPlugin
class _InstallWheelPlugin(WorkerPlugin):
name = f"install-{package_name}"
def __init__(
self,
wheel_bytes,
wheel_filename,
package_name,
expected_version,
env_installed_wheel,
env_expected_version,
):
self.wheel_bytes = wheel_bytes
self.wheel_filename = wheel_filename
self.package_name = package_name
self.expected_version = expected_version
self.env_installed_wheel = env_installed_wheel
self.env_expected_version = env_expected_version
def setup(self, worker):
return _install_wheel_bytes_embedded(
self.wheel_bytes,
self.wheel_filename,
self.package_name,
self.expected_version,
self.env_installed_wheel,
self.env_expected_version,
)
# Install on the scheduler first using the embedded function in this
# scope (cloudpickle serialises it by value). The scheduler process
# isn't a "worker" so the plugin mechanism above doesn't cover it.
sched_msg = client.run_on_scheduler(
_install_wheel_bytes_embedded,
wheel_bytes,
wheel_filename,
package_name,
expected_version,
env_installed_wheel,
env_expected_version,
)
logger.info(f" scheduler: {sched_msg}")
# Install on all current workers and register for all future ones.
plugin = _InstallWheelPlugin(
wheel_bytes,
wheel_filename,
package_name,
expected_version,
env_installed_wheel,
env_expected_version,
)
results = client.register_plugin(plugin, name=plugin.name)
for worker_addr, msg in results.items():
logger.info(f" {worker_addr}: {msg}")
worker_addresses = list(client.scheduler_info()["workers"].keys())
fingerprints = client.run(
_worker_fingerprint_embedded, package_name, workers=worker_addresses
)
for worker_addr, fp in fingerprints.items():
logger.info(
" %s version=%s file=%s pid=%s",
worker_addr,
fp.get("version"),
fp.get("module_file"),
fp.get("pid"),
)
if fp.get("version") != expected_version:
raise RuntimeError(
f"Worker {worker_addr} version mismatch: "
f"expected {expected_version}, got {fp.get('version')}"
)
logger.info(
f"Scheduler and all workers have {package_name} installed and "
"verified at runtime."
)
def _register_package(client) -> None:
"""Build pyfracval's own wheel and install it on all workers.
Thin wrapper around :func:`install_wheel_on_workers` -- kept for
backwards compatibility and as pyfracval's own "build from local source
first" use case; other packages (e.g. a compiled extension like
pyfastmm, which needs a wheel built for the *worker's* platform, not
whatever this machine happens to be) call
:func:`install_wheel_on_workers` directly with an already-built wheel
instead of going through this function.
"""
wheel_path = _build_wheel()
expected_version = _project_version()
install_wheel_on_workers(client, wheel_path, "pyfracval", expected_version)
[docs]
def get_client(
scheduler_address: str | None = None,
n_workers: int | None = None,
install_package: bool = False,
):
"""Return a Dask distributed Client.
If *scheduler_address* is given, connect to a running scheduler at that
address (e.g. ``"tcp://host:8786"``). Otherwise start a local
``LocalCluster`` with *n_workers* workers (defaults to the number of
CPU cores when *n_workers* is ``None``).
When *install_package* is ``True`` **and** a remote scheduler is used, the
local ``pyfracval`` package is built into a wheel and installed on all
workers via a ``WorkerPlugin`` before the client is returned. This is
required whenever the workers do not have ``pyfracval`` pre-installed
(e.g. a generic Dask Docker image).
Parameters
----------
scheduler_address:
Address of a remote Dask scheduler. ``None`` → use a local cluster.
n_workers:
Number of workers for a local cluster. Ignored when connecting to a
remote scheduler.
install_package:
When ``True`` and using a remote scheduler, build + install
``pyfracval`` on all workers before returning.
Returns
-------
dask.distributed.Client
"""
import dask
from dask.distributed import Client, LocalCluster # lazy import
if scheduler_address is not None:
logger.info(f"Connecting to remote Dask scheduler at {scheduler_address}")
# register_plugin() (used by install_wheel_on_workers() below, and
# called up to 3x per campaign launch -- pyfracval/spcwth/pyfastmm)
# is a single blocking RPC that waits for every currently-connected
# worker to finish a real `pip install` before the scheduler
# replies at all -- the socket sits silent the whole time. Dask's
# default distributed.comm.timeouts.tcp is only 30s, comfortably
# exceeded once a cluster has ~20+ workers all installing
# concurrently (observed in practice: CommClosedError("Stream is
# closed") mid-registration, not a real network failure -- just
# this timeout firing while the install was still legitimately in
# progress). Set once, globally, before the client exists; cheap
# insurance for the rest of this process's Dask RPCs too.
dask.config.set(
{
"distributed.comm.timeouts.tcp": "120s",
"distributed.comm.timeouts.connect": "60s",
}
)
client = Client(scheduler_address)
if install_package:
_register_package(client)
return client
if n_workers is None:
# LocalCluster(n_workers=None) does NOT simply use all CPU cores -
# Dask's own default heuristic also factors in currently-available
# system memory and can pick far fewer workers than cores on a
# machine with other things running (observed: 4 workers on a
# 16-core/64GB desktop with a few GB already committed to unrelated
# apps). Resolve explicitly to the actual core count so a
# compute-bound batch workload (a parameter sweep) gets what the
# docstring above promises, rather than a memory-conservative guess.
n_workers = os.cpu_count() or 1
logger.info(f"Starting local Dask cluster with n_workers={n_workers!r}")
cluster = LocalCluster(n_workers=n_workers)
return Client(cluster)