From 6c11eaf337fb8f565a17f90b20108b3fe27951a7 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 9 Oct 2026 09:14:07 -0400 Subject: [PATCH] Keep the telemetry emitter service off the modelling stack The emitter service is launched as `python -m microcosm.build.telemetry_emitter_service`, so Python ran the eager `microcosm/build/__init__.py` first and the service imported torch, pandas, numpy, scipy and pyarrow it never uses. On a loaded build host that took 5-9 s against a 3 s readiness budget, and the build continued without hosted telemetry. - Load `microcosm.build`'s public names on first use (PEP 562). `__all__` is unchanged; the original imports stay under TYPE_CHECKING for type checkers and the AST identity closures. - Read the microcosm-frame version for the compatibility gate from distribution metadata, falling back to the module attribute when frame is not installed, so the gate no longer imports the kernel. - Import huggingface_hub in the collector only when the delivery worker needs a cached token. - Wait up to 30 s for the service; start() still returns as soon as the service answers or exits. Report an exit with its status separately from a timeout, and name the wait in the timeout warning. - Star-import `microcosm.build` in the Modal environment and runner smoke probes so they still prove the modelling stack imports. Measured at load 78-162: service ready in 0.41-0.84 s, down from 5.3-9.2 s. Co-Authored-By: Claude Opus 5.5 --- DESIGN.md | 7 +- .../telemetry-emitter-startup.fixed.md | 1 + .../src/microcosm/build/__init__.py | 480 ++++++++++++------ .../src/microcosm/build/telemetry_emitter.py | 11 +- .../build/telemetry_emitter_constants.py | 17 +- .../telemetry_emitter_service/collector.py | 13 +- .../tests/engine_free/shared/test_compat.py | 97 ++++ .../engine_free/shared/test_package_init.py | 139 +++++ .../test_telemetry_emitter_service_imports.py | 103 ++++ .../shared/test_telemetry_emitter_startup.py | 209 ++++++++ .../us/test_us_modal_stage_plan_tool.py | 2 +- tools/modal_us_stage.py | 4 +- tools/modal_us_stage_plan.py | 4 +- 13 files changed, 929 insertions(+), 158 deletions(-) create mode 100644 changelog.d/telemetry-emitter-startup.fixed.md create mode 100644 packages/microcosm-build/tests/engine_free/shared/test_compat.py create mode 100644 packages/microcosm-build/tests/engine_free/shared/test_package_init.py create mode 100644 packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter_service_imports.py create mode 100644 packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter_startup.py diff --git a/DESIGN.md b/DESIGN.md index 19b23bf4a0..9c3d4d6d2a 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -256,8 +256,11 @@ never as namespace squatters. **Constellation versioning has a mechanism, not just an intent.** Each shard pins `microcosm-frame>=X, str: + """Return the installed microcosm-frame version without importing the kernel. + + Importing :mod:`microcosm.frame` loads numpy and pandas, which every process + importing any build submodule would then pay for. The gate guards against a + resolver assembling an incompatible installed set, and the distribution + metadata is the version that was installed. A frame importable from the path + without an installed distribution falls back to its module attribute. + """ + + try: + return _metadata.version("microcosm-frame") + except _metadata.PackageNotFoundError: + from microcosm.frame import __version__ + + return __version__ + + def _assert_frame_compatible(version: str, required: tuple[int, int]) -> None: """Raise unless the installed microcosm-frame is the expected series.""" parts = version.split(".") @@ -47,151 +73,315 @@ def _assert_frame_compatible(version: str, required: tuple[int, int]) -> None: ) -_assert_frame_compatible(_frame_version, _REQUIRED_FRAME_SERIES) +_assert_frame_compatible(_installed_frame_version(), _REQUIRED_FRAME_SERIES) + +if TYPE_CHECKING: + from microcosm.build.chronicle_epoch import ( + ACCEPTED_CONSUMER_ARTIFACT_SCHEMA_VERSIONS, + ACCEPTED_CONSUMER_FACT_SCHEMA_VERSIONS, + CHRONICLE_EPOCH, + DECLARED_IDENTITIES, + DECLARED_IDENTITY_EPOCHS, + EPOCHS, + LEDGER_EPOCH, + UNDECLARED, + fact_key_epoch, + fact_key_epoch_label, + feed_fact_key_epochs, + feed_undeclared_fact_key_domains, + is_chronicle_fact_key, + parse_fact_key, + ) + from microcosm.build.country_spec import ( + CountryResourceRow, + CountrySpec, + GateSelectionSpec, + GatesManifest, + GeographySpineManifest, + GeographySpineSpec, + ReleaseContractManifest, + ResolvedCountrySpec, + country_stage_plan, + load_country_spec, + ) + from microcosm.build.gate_battery import ( + BlockingMode, + EvidenceContext, + FunctionBinding, + GateBatteryBlockedError, + GateBatteryRun, + GateBinding, + GateOutcome, + GatePhaseReport, + GateStatus, + evaluate_phase, + validate_gate_parameters, + ) + from microcosm.build.gates import ( + FitWeightRecord, + GateReport, + GateResult, + TargetCoverageRequirement, + TargetFitRequirement, + aggregate_admin_gate, + area_support_gate, + column_implication_gate, + default_valued_columns_gate, + enum_domain_gate, + export_surface_gate, + exported_nonzero_gate, + formula_owned_export_gate, + input_column_coverage_gate, + input_mass_parity_gate, + ledger_compile_parity_gate, + ledger_compile_parity_signed_differences, + macro_realism_gate, + nonconstant_columns_gate, + nonnegative_columns_gate, + parity_gate, + per_family_fit_gate, + relative_error_loss, + source_coverage_gate, + source_stage_input_coverage_gate, + support_gate, + tail_concentration_gate, + target_fit_gate, + target_profile_coverage_gate, + target_surface_gate, + weights_audit_gate, + ) + from microcosm.build.holdout import ( + hash_holdout_uniform, + hash_holdout_unit, + rotated_folds, + summarize_rotations, + ) + from microcosm.build.ledger_artifact import ( + LedgerConsumerArtifact, + add_ledger_artifact_args, + load_ledger_consumer_artifact, + resolve_ledger_artifact, + ) + from microcosm.build.ledger_targets import ( + LedgerTargetMapping, + LedgerTargetSelection, + UnsupportedLedgerTarget, + apply_ledger_target_profile, + select_ledger_targets, + select_ledger_targets_from_jsonl, + target_spec_from_ledger_fact, + ) + + # Only ``logbook_env_names`` is re-exported here. The reader function is + # spelled ``logbook_env`` — the same name as its module — so binding it on + # the package would make ``import microcosm.build.logbook_env as env`` hand + # back the function instead of the module. Import the function from its + # module: ``from microcosm.build.logbook_env import logbook_env``. + from microcosm.build.logbook_env import ( + logbook_env_names, + ) + from microcosm.build.monetary_profile import ( + MonetaryTargetContract, + MonetaryTargetProfile, + ) + from microcosm.build.monetary_targets import ( + MonetaryBasis, + PreparedMonetaryMeasure, + bind_monetary_target, + prepare_monetary_measure, + ) + from microcosm.build.observation import ( + ObservedTransform, + StageEventObserver, + StageEventRun, + StageEventStatus, + StageObservation, + StageObservationRun, + StageObserver, + ) + from microcosm.build.plan import ( + DonorSpec, + Stage, + StagePlan, + StageRecord, + ) + from microcosm.build.source_runtime import ( + SourceRuntimeConfig, + SourceRuntimeContext, + SourceRuntimeError, + UnsupportedSourceOperationError, + run_source_stage, + ) + from microcosm.build.staging import ( + DEFAULT_STAGING_PREFIX, + LATEST_STAGING_POINTER, + RUNS_INDEX, + STAGING_SCHEMA_VERSION, + StagingRunBundleWriter, + ) + + +#: Each public name, grouped by the submodule that defines it. The +#: ``TYPE_CHECKING`` imports above name the same pairs for type checkers and +#: for the source-import closures behind gate-binding and worker identities. +_EXPORTS_BY_MODULE: dict[str, tuple[str, ...]] = { + "chronicle_epoch": ( + "ACCEPTED_CONSUMER_ARTIFACT_SCHEMA_VERSIONS", + "ACCEPTED_CONSUMER_FACT_SCHEMA_VERSIONS", + "CHRONICLE_EPOCH", + "DECLARED_IDENTITIES", + "DECLARED_IDENTITY_EPOCHS", + "EPOCHS", + "LEDGER_EPOCH", + "UNDECLARED", + "fact_key_epoch", + "fact_key_epoch_label", + "feed_fact_key_epochs", + "feed_undeclared_fact_key_domains", + "is_chronicle_fact_key", + "parse_fact_key", + ), + "country_spec": ( + "CountryResourceRow", + "CountrySpec", + "GateSelectionSpec", + "GatesManifest", + "GeographySpineManifest", + "GeographySpineSpec", + "ReleaseContractManifest", + "ResolvedCountrySpec", + "country_stage_plan", + "load_country_spec", + ), + "gate_battery": ( + "BlockingMode", + "EvidenceContext", + "FunctionBinding", + "GateBatteryBlockedError", + "GateBatteryRun", + "GateBinding", + "GateOutcome", + "GatePhaseReport", + "GateStatus", + "evaluate_phase", + "validate_gate_parameters", + ), + "gates": ( + "FitWeightRecord", + "GateReport", + "GateResult", + "TargetCoverageRequirement", + "TargetFitRequirement", + "aggregate_admin_gate", + "area_support_gate", + "column_implication_gate", + "default_valued_columns_gate", + "enum_domain_gate", + "export_surface_gate", + "exported_nonzero_gate", + "formula_owned_export_gate", + "input_column_coverage_gate", + "input_mass_parity_gate", + "ledger_compile_parity_gate", + "ledger_compile_parity_signed_differences", + "macro_realism_gate", + "nonconstant_columns_gate", + "nonnegative_columns_gate", + "parity_gate", + "per_family_fit_gate", + "relative_error_loss", + "source_coverage_gate", + "source_stage_input_coverage_gate", + "support_gate", + "tail_concentration_gate", + "target_fit_gate", + "target_profile_coverage_gate", + "target_surface_gate", + "weights_audit_gate", + ), + "holdout": ( + "hash_holdout_uniform", + "hash_holdout_unit", + "rotated_folds", + "summarize_rotations", + ), + "ledger_artifact": ( + "LedgerConsumerArtifact", + "add_ledger_artifact_args", + "load_ledger_consumer_artifact", + "resolve_ledger_artifact", + ), + "ledger_targets": ( + "LedgerTargetMapping", + "LedgerTargetSelection", + "UnsupportedLedgerTarget", + "apply_ledger_target_profile", + "select_ledger_targets", + "select_ledger_targets_from_jsonl", + "target_spec_from_ledger_fact", + ), + # Never ``logbook_env`` itself: see the note on its import above. + "logbook_env": ("logbook_env_names",), + "monetary_profile": ( + "MonetaryTargetContract", + "MonetaryTargetProfile", + ), + "monetary_targets": ( + "MonetaryBasis", + "PreparedMonetaryMeasure", + "bind_monetary_target", + "prepare_monetary_measure", + ), + "observation": ( + "ObservedTransform", + "StageEventObserver", + "StageEventRun", + "StageEventStatus", + "StageObservation", + "StageObservationRun", + "StageObserver", + ), + "plan": ( + "DonorSpec", + "Stage", + "StagePlan", + "StageRecord", + ), + "source_runtime": ( + "SourceRuntimeConfig", + "SourceRuntimeContext", + "SourceRuntimeError", + "UnsupportedSourceOperationError", + "run_source_stage", + ), + "staging": ( + "DEFAULT_STAGING_PREFIX", + "LATEST_STAGING_POINTER", + "RUNS_INDEX", + "STAGING_SCHEMA_VERSION", + "StagingRunBundleWriter", + ), +} +_EXPORTS = { + name: module for module, names in _EXPORTS_BY_MODULE.items() for name in names +} + + +def __getattr__(name: str) -> object: + """Load a public name from its defining submodule on first use.""" + + try: + module = _EXPORTS[name] + except KeyError: + # AttributeError, not ImportError: ``from microcosm.build import + # `` falls back to importing the submodule only on it. + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + return getattr(_import_module(f"{__name__}.{module}"), name) + -from microcosm.build.chronicle_epoch import ( # noqa: E402 - after the compat gate - ACCEPTED_CONSUMER_ARTIFACT_SCHEMA_VERSIONS, - ACCEPTED_CONSUMER_FACT_SCHEMA_VERSIONS, - CHRONICLE_EPOCH, - DECLARED_IDENTITIES, - DECLARED_IDENTITY_EPOCHS, - EPOCHS, - LEDGER_EPOCH, - UNDECLARED, - fact_key_epoch, - fact_key_epoch_label, - feed_fact_key_epochs, - feed_undeclared_fact_key_domains, - is_chronicle_fact_key, - parse_fact_key, -) -from microcosm.build.country_spec import ( # noqa: E402 - after the compat gate - CountryResourceRow, - CountrySpec, - GateSelectionSpec, - GatesManifest, - GeographySpineManifest, - GeographySpineSpec, - ReleaseContractManifest, - ResolvedCountrySpec, - country_stage_plan, - load_country_spec, -) -from microcosm.build.gate_battery import ( # noqa: E402 - after the compat gate - BlockingMode, - EvidenceContext, - FunctionBinding, - GateBatteryBlockedError, - GateBatteryRun, - GateBinding, - GateOutcome, - GatePhaseReport, - GateStatus, - evaluate_phase, - validate_gate_parameters, -) -from microcosm.build.gates import ( # noqa: E402 - after the compat gate - FitWeightRecord, - GateReport, - GateResult, - TargetCoverageRequirement, - TargetFitRequirement, - aggregate_admin_gate, - area_support_gate, - column_implication_gate, - default_valued_columns_gate, - enum_domain_gate, - export_surface_gate, - exported_nonzero_gate, - formula_owned_export_gate, - input_column_coverage_gate, - input_mass_parity_gate, - ledger_compile_parity_gate, - ledger_compile_parity_signed_differences, - macro_realism_gate, - nonconstant_columns_gate, - nonnegative_columns_gate, - parity_gate, - per_family_fit_gate, - relative_error_loss, - source_coverage_gate, - source_stage_input_coverage_gate, - support_gate, - tail_concentration_gate, - target_fit_gate, - target_profile_coverage_gate, - target_surface_gate, - weights_audit_gate, -) -from microcosm.build.holdout import ( # noqa: E402 - after the compat gate - hash_holdout_uniform, - hash_holdout_unit, - rotated_folds, - summarize_rotations, -) -from microcosm.build.ledger_artifact import ( # noqa: E402 - after the compat gate - LedgerConsumerArtifact, - add_ledger_artifact_args, - load_ledger_consumer_artifact, - resolve_ledger_artifact, -) -from microcosm.build.ledger_targets import ( # noqa: E402 - after the compat gate - LedgerTargetMapping, - LedgerTargetSelection, - UnsupportedLedgerTarget, - apply_ledger_target_profile, - select_ledger_targets, - select_ledger_targets_from_jsonl, - target_spec_from_ledger_fact, -) +def __dir__() -> list[str]: + return sorted(set(globals()) | set(__all__)) -# Only ``logbook_env_names`` is re-exported here. The reader function is -# spelled ``logbook_env`` — the same name as its module — so binding it on -# the package would make ``import microcosm.build.logbook_env as env`` hand -# back the function instead of the module. Import the function from its -# module: ``from microcosm.build.logbook_env import logbook_env``. -from microcosm.build.logbook_env import ( # noqa: E402 - after the compat gate - logbook_env_names, -) -from microcosm.build.monetary_profile import ( # noqa: E402 - after compat gate - MonetaryTargetContract, - MonetaryTargetProfile, -) -from microcosm.build.monetary_targets import ( # noqa: E402 - after compat gate - MonetaryBasis, - PreparedMonetaryMeasure, - bind_monetary_target, - prepare_monetary_measure, -) -from microcosm.build.observation import ( # noqa: E402 - after the compat gate - ObservedTransform, - StageEventObserver, - StageEventRun, - StageEventStatus, - StageObservation, - StageObservationRun, - StageObserver, -) -from microcosm.build.plan import ( # noqa: E402 - after the compat gate - DonorSpec, - Stage, - StagePlan, - StageRecord, -) -from microcosm.build.source_runtime import ( # noqa: E402 - after the compat gate - SourceRuntimeConfig, - SourceRuntimeContext, - SourceRuntimeError, - UnsupportedSourceOperationError, - run_source_stage, -) -from microcosm.build.staging import ( # noqa: E402 - after the compat gate - DEFAULT_STAGING_PREFIX, - LATEST_STAGING_POINTER, - RUNS_INDEX, - STAGING_SCHEMA_VERSION, - StagingRunBundleWriter, -) __version__ = "0.1.0" diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py index 761f49617f..b6683846b6 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py @@ -29,6 +29,7 @@ NO_LOCAL_SOCKET_WARNING, QUEUE_WARNING, RUNTIME_DIRECTORY_PREFIX, + SERVICE_EXITED_WARNING, SERVICE_NOT_READY_WARNING, SERVICE_PROCESS_EXIT_TIMEOUT_SECONDS, SERVICE_READY_TIMEOUT_SECONDS, @@ -233,7 +234,15 @@ def start( file=sys.stderr, ) else: - print(SERVICE_NOT_READY_WARNING, file=sys.stderr) + returncode = process.poll() + print( + SERVICE_NOT_READY_WARNING.format( + timeout_seconds=startup_timeout_seconds + ) + if returncode is None + else SERVICE_EXITED_WARNING.format(returncode=returncode), + file=sys.stderr, + ) if process is not None: _terminate_process(process) try: diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_constants.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_constants.py index 57a477f82b..48909e8bb7 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_constants.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_constants.py @@ -14,7 +14,12 @@ DEFAULT_SEND_TIMEOUT_SECONDS: Final = 0.2 DEFAULT_HEARTBEAT_SECONDS: Final = 60.0 -DEFAULT_STARTUP_TIMEOUT_SECONDS: Final = 3.0 +# start() returns as soon as the service answers its first ping or exits, so +# only a live service that never answers costs the whole wait, while giving up +# early discards hosted telemetry for the entire build. Startup is interpreter +# start plus the SQLAlchemy spool: 0.4-0.8 s on a build host at load average +# 80-160 on 2026-10-09, where single imports also stalled for up to 24 s. +DEFAULT_STARTUP_TIMEOUT_SECONDS: Final = 30.0 STARTUP_POLL_SECONDS: Final = 0.02 SERVICE_READY_TIMEOUT_SECONDS: Final = 0.1 SERVICE_PROCESS_EXIT_TIMEOUT_SECONDS: Final = 1.0 @@ -27,9 +32,15 @@ "warning: the local telemetry emitter service could not start: " "{error_type}: {error}" ) +SERVICE_EXITED_WARNING: Final = ( + "warning: the local telemetry emitter service exited with status " + "{returncode} before it became ready; the build will continue without " + "hosted telemetry." +) SERVICE_NOT_READY_WARNING: Final = ( - "warning: the local telemetry emitter service did not become ready; " - "the build will continue without hosted telemetry." + "warning: the local telemetry emitter service did not become ready " + "within {timeout_seconds:g} s; the build will continue without hosted " + "telemetry." ) QUEUE_WARNING: Final = ( "warning: the local telemetry emitter service could not queue an update " diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py index 795c61a177..be75f308c9 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py @@ -13,8 +13,6 @@ from typing import Any from urllib.parse import urlsplit -from huggingface_hub import get_token - from microcosm.build.telemetry_emitter_service.constants import ( COLLECTOR_RESPONSE_TOO_LARGE_ERROR, COLLECTOR_URL_HTTPS_ERROR, @@ -116,11 +114,18 @@ def _http_post( def _huggingface_token() -> str | None: - return ( + environment_token = ( os.environ.get("HF_TOKEN", "").strip() or os.environ.get("HUGGINGFACE_TOKEN", "").strip() - or get_token() ) + if environment_token: + return environment_token + # Imported here, not at module level: huggingface_hub brings httpx, yaml + # and asyncio, and only the delivery worker needs it, after the service + # has answered the build's readiness ping. + from huggingface_hub import get_token + + return get_token() def _registration_key(registration: Mapping[str, Any]) -> str: diff --git a/packages/microcosm-build/tests/engine_free/shared/test_compat.py b/packages/microcosm-build/tests/engine_free/shared/test_compat.py new file mode 100644 index 0000000000..ff9049b245 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_compat.py @@ -0,0 +1,97 @@ +"""The constellation mechanism: the kernel-compatibility gate at import. + +DESIGN.md requires each shard to assert kernel compatibility at import so a +resolver that ignores ``[tool.uv.sources]`` cannot silently assemble an +incompatible pair. microcosm-build reads the installed frame version from +distribution metadata, so the gate does not load the kernel's numpy and pandas +stack. These tests exercise the gate directly (the real import already ran it +successfully, or the suite would not have loaded). +""" + +from __future__ import annotations + +from importlib import metadata + +import pytest +from hypothesis import given +from hypothesis import strategies as st + +import microcosm.build +import microcosm.frame +from microcosm.build import _assert_frame_compatible, _installed_frame_version + + +def test_build_declares_its_own_version() -> None: + """The shard exposes its version for the constellation matrix.""" + assert microcosm.build.__version__ == "0.1.0" + + +def test_gate_reads_the_installed_frame_distribution() -> None: + """The gate's version is the installed distribution's version.""" + assert _installed_frame_version() == metadata.version("microcosm-frame") + + +def test_installed_distribution_matches_the_kernel_module() -> None: + """Metadata and ``frame.__version__`` are hand-written separately; keep them equal. + + The gate reads metadata while microcosm-fit and microcosm-calibrate read the + module attribute, so both sources must name the same kernel. + """ + assert metadata.version("microcosm-frame") == microcosm.frame.__version__ + + +def test_gate_falls_back_to_the_module_without_a_distribution(monkeypatch) -> None: + """A frame importable from the path but not installed still gets checked.""" + + def not_installed(name: str) -> str: + raise metadata.PackageNotFoundError(name) + + monkeypatch.setattr(microcosm.build._metadata, "version", not_installed) + assert _installed_frame_version() == microcosm.frame.__version__ + + +def test_compat_gate_accepts_the_matching_series() -> None: + """The installed kernel passes the gate (this is the live configuration).""" + _assert_frame_compatible("0.1.0", (0, 1)) + _assert_frame_compatible("0.1.5", (0, 1)) + + +def test_compat_gate_rejects_a_too_old_or_too_new_kernel() -> None: + """A kernel outside the required 0.x minor series is refused at import.""" + with pytest.raises(ImportError, match="requires microcosm-frame 0.1.x"): + _assert_frame_compatible("0.0.9", (0, 1)) + with pytest.raises(ImportError, match="0.2.0 is installed"): + _assert_frame_compatible("0.2.0", (0, 1)) + + +def test_compat_gate_uses_major_only_from_1_0() -> None: + """From 1.0 on, the gate matches the major and tolerates any minor.""" + _assert_frame_compatible("1.4.2", (1, 0)) + with pytest.raises(ImportError, match="requires microcosm-frame 2.x"): + _assert_frame_compatible("1.9.9", (2, 0)) + + +def test_compat_gate_rejects_an_unparseable_version() -> None: + """A version string the gate cannot parse is a clear ImportError.""" + with pytest.raises(ImportError, match="cannot parse"): + _assert_frame_compatible("not-a-version", (0, 1)) + + +_PART = st.integers(min_value=0, max_value=30) + + +@given(major=_PART, minor=_PART, patch=_PART, required_major=_PART, required=_PART) +def test_compat_gate_accepts_exactly_the_required_series( + major: int, minor: int, patch: int, required_major: int, required: int +) -> None: + """Pre-1.0 the gate pins the minor series; from 1.0 it pins the major only.""" + version = f"{major}.{minor}.{patch}" + series = (required_major, required) + accepted = ( + (major, minor) == series if required_major == 0 else major == required_major + ) + if accepted: + _assert_frame_compatible(version, series) + else: + with pytest.raises(ImportError, match=f"{version} is installed"): + _assert_frame_compatible(version, series) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_package_init.py b/packages/microcosm-build/tests/engine_free/shared/test_package_init.py new file mode 100644 index 0000000000..23ba399c56 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_package_init.py @@ -0,0 +1,139 @@ +"""The build package loads its public names lazily (PEP 562). + +``import microcosm.build`` runs only the frame compatibility gate. Each name in +``__all__`` loads from its defining submodule on first use, so a process that +needs one submodule, such as the telemetry emitter service, does not pay for +torch, pandas and the calibration stack. +""" + +from __future__ import annotations + +import ast +import json +import subprocess +import sys +import types +from importlib import import_module +from pathlib import Path + +import pytest + +import microcosm.build +from microcosm.build import _EXPORTS, _EXPORTS_BY_MODULE + +#: Third-party and kernel modules that only the re-exported submodules load. +_MODELLING_STACK = ( + "torch", + "scipy", + "pandas", + "numpy", + "pyarrow", + "microcosm.frame", + "microcosm.calibrate", +) + + +def _loaded(modules: set[str], roots: tuple[str, ...]) -> list[str]: + return sorted( + name + for name in modules + if any(name == root or name.startswith(f"{root}.") for root in roots) + ) + + +def test_export_map_names_every_public_name_once() -> None: + """``__all__`` is the export map plus ``__version__``, with no name twice.""" + names = [name for names in _EXPORTS_BY_MODULE.values() for name in names] + assert len(names) == len(set(names)) == len(_EXPORTS) + assert set(_EXPORTS) | {"__version__"} == set(microcosm.build.__all__) + assert len(microcosm.build.__all__) == len(set(microcosm.build.__all__)) + + +def test_every_public_name_is_its_defining_modules_object() -> None: + """Lazy resolution returns the same object the eager import used to bind.""" + for name, module in _EXPORTS.items(): + defining = import_module(f"microcosm.build.{module}") + assert getattr(microcosm.build, name) is getattr(defining, name), name + + +def test_star_import_binds_every_public_name() -> None: + namespace: dict[str, object] = {} + exec("from microcosm.build import *", namespace) + assert set(microcosm.build.__all__) <= set(namespace) + + +def test_type_checking_imports_match_the_export_map() -> None: + """The static import block and the runtime map name the same pairs. + + Gate-binding and worker identities hash the source-import closure, which + reads the ``TYPE_CHECKING`` block; the runtime reads ``_EXPORTS_BY_MODULE``. + """ + tree = ast.parse(Path(microcosm.build.__file__).read_text()) + blocks = [ + node + for node in tree.body + if isinstance(node, ast.If) + and isinstance(node.test, ast.Name) + and node.test.id == "TYPE_CHECKING" + ] + assert len(blocks) == 1 + static = { + node.module.removeprefix("microcosm.build."): tuple( + alias.name for alias in node.names + ) + for node in blocks[0].body + if isinstance(node, ast.ImportFrom) + } + assert static == _EXPORTS_BY_MODULE + + +def test_logbook_env_attribute_stays_the_module() -> None: + """Only ``logbook_env_names`` is exported; the name stays the submodule.""" + import microcosm.build.logbook_env + + assert "logbook_env" not in _EXPORTS + assert isinstance(microcosm.build.logbook_env, types.ModuleType) + + +def test_unknown_names_raise_attribute_error() -> None: + """Lookups that miss raise AttributeError, as ``from`` imports require.""" + with pytest.raises(AttributeError, match="no_such_export"): + microcosm.build.no_such_export # noqa: B018 + assert not hasattr(microcosm.build, "__wrapped__") + assert getattr(microcosm.build, "no_such_export", None) is None + + +def test_dir_lists_every_public_name() -> None: + assert set(microcosm.build.__all__) <= set(dir(microcosm.build)) + + +_PACKAGE_IMPORT_PROBE = r""" +import json +import sys + +import microcosm.build +from microcosm.build import CHRONICLE_EPOCH, chronicle_epoch + +assert chronicle_epoch.__name__ == "microcosm.build.chronicle_epoch" +assert CHRONICLE_EPOCH is chronicle_epoch.CHRONICLE_EPOCH +print(json.dumps(sorted(sys.modules))) +""" + + +def test_package_import_leaves_the_modelling_stack_unloaded() -> None: + """A fresh interpreter importing the package loads no modelling stack. + + ``chronicle_epoch`` imports only the standard library, so loading one of its + exports, or the submodule by name, must not pull the other exports in. + """ + completed = subprocess.run( + [sys.executable, "-c", _PACKAGE_IMPORT_PROBE], + capture_output=True, + text=True, + timeout=300, + check=False, + ) + assert completed.returncode == 0, completed.stderr + loaded = set(json.loads(completed.stdout.splitlines()[-1])) + assert not _loaded(loaded, _MODELLING_STACK) + assert not _loaded(loaded, ("microcosm.build.gates", "microcosm.build.plan")) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter_service_imports.py b/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter_service_imports.py new file mode 100644 index 0000000000..ac3679a116 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter_service_imports.py @@ -0,0 +1,103 @@ +"""The emitter service's readiness path stays off the modelling stack. + +``LocalTelemetryEmitter.start`` waits for the service to answer a ping before +the build proceeds, and abandons hosted telemetry if it does not. Everything +the service imports before that answer is startup latency, so the probe below +does the service's pre-readiness work in a fresh interpreter and checks which +modules it loaded. The check is on the module set, not wall time, so it holds +on a loaded host. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys + +from microcosm.build.telemetry_emitter_service import collector as collector_module + +#: Modules the service never needs: the modelling stack the build package's +#: exports pull in, the other shards, and the Hugging Face client (needed only +#: by the delivery worker after the service answers, and only without an +#: environment token). +_FORBIDDEN = ( + "torch", + "scipy", + "pandas", + "numpy", + "pyarrow", + "microcosm.frame", + "microcosm.calibrate", + "microcosm.fit", + "microcosm.graph", + "microcosm.diagnostics", + "microcosm.data", + "microcosm.build.country_spec", + "microcosm.build.gates", + "microcosm.build.ledger_targets", + "huggingface_hub", + "httpx", +) + +# Same imports as ``python -m microcosm.build.telemetry_emitter_service`` (the +# package and ``main``), then the work ``main`` and ``EmitterService.run`` do +# before binding the socket, then the token lookup with an environment token. +_READINESS_PROBE = r""" +import json +import os +import sys + +import microcosm.build.telemetry_emitter_service.main +from microcosm.build.telemetry_emitter_service import collector +from microcosm.build.telemetry_emitter_service.collector import CollectorDelivery +from microcosm.build.telemetry_emitter_service.resources import ProcessTreeSampler +from microcosm.build.telemetry_emitter_service.spool import EventSpool + +spool = EventSpool(sys.argv[1]) +CollectorDelivery(spool, development_collector_url="http://127.0.0.1:1") +ProcessTreeSampler(os.getpid()) +spool.register({"run_id": "probe", "producer_id": "probe"}) +print(json.dumps({"token": collector._huggingface_token(), "modules": sorted(sys.modules)})) +""" + + +def test_service_readiness_path_loads_no_modelling_stack(tmp_path) -> None: + environment = {**os.environ, "HF_TOKEN": "hf-environment-token"} + completed = subprocess.run( + [sys.executable, "-c", _READINESS_PROBE, str(tmp_path / "events.sqlite3")], + capture_output=True, + text=True, + env=environment, + timeout=300, + check=False, + ) + assert completed.returncode == 0, completed.stderr + probe = json.loads(completed.stdout.splitlines()[-1]) + assert probe["token"] == "hf-environment-token" + leaked = sorted( + name + for name in probe["modules"] + if any(name == root or name.startswith(f"{root}.") for root in _FORBIDDEN) + ) + assert not leaked + + +def test_token_lookup_falls_back_to_the_cached_hugging_face_login() -> None: + """Without an environment token the deferred import reads the cached login. + + The autouse telemetry fixture clears the token variables and points the + Hugging Face token path into this test's temporary directory. + """ + from huggingface_hub import constants + + token_path = constants.HF_TOKEN_PATH + os.makedirs(os.path.dirname(token_path), exist_ok=True) + with open(token_path, "w") as handle: + handle.write("hf-cached-login\n") + + assert collector_module._huggingface_token() == "hf-cached-login" + + +def test_token_lookup_is_empty_without_any_credential() -> None: + assert not collector_module._huggingface_token() diff --git a/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter_startup.py b/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter_startup.py new file mode 100644 index 0000000000..537d5ac0e1 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter_startup.py @@ -0,0 +1,209 @@ +"""How long ``LocalTelemetryEmitter.start`` waits for its service, and what it keeps. + +A fake clock drives the real startup loop. The service process, its readiness +ping and ``time`` are fakes, so each example runs in microseconds whatever the +host load, and the wait's contract is asserted exactly: + +* a service that answers before the budget runs out is kept, however slowly + it started; +* a service that exits first is reported with its exit status, at once; +* a live service that never answers is abandoned at the budget, no later than + one poll after it; +* nothing is sent to a service that never answered. +""" + +from __future__ import annotations + +import itertools +import math +from types import SimpleNamespace + +import pytest +from hypothesis import HealthCheck, assume, given, settings +from hypothesis import strategies as st + +from microcosm.build import telemetry_emitter as module +from microcosm.build.telemetry_emitter import LocalTelemetryEmitter +from microcosm.build.telemetry_emitter_constants import ( + DEFAULT_STARTUP_TIMEOUT_SECONDS, + SOCKET_FILENAME, + STARTUP_POLL_SECONDS, +) +from microcosm.build.telemetry_protocol import BUILD_STARTED_MESSAGE + +_RUNTIME_DIRECTORIES = itertools.count() +_EXIT_STATUS = 3 +# Event times this far apart cannot share a poll, even with float drift in the +# fake clock, so each example has one unambiguous expected outcome. +_SEPARATION = 3 * STARTUP_POLL_SECONDS + + +class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def monotonic(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.now += seconds + + +class _ServiceProcess: + """A service process that exits at ``exit_at`` on the fake clock.""" + + def __init__(self, clock: _Clock, exit_at: float | None) -> None: + self._clock = clock + self._exit_at = exit_at + self.terminated = False + self.waited = False + + def poll(self) -> int | None: + if self.terminated: + return -15 + if self._exit_at is not None and self._clock.now >= self._exit_at: + return _EXIT_STATUS + return None + + def terminate(self) -> None: + self.terminated = True + + def kill(self) -> None: + self.terminated = True + + def wait(self, timeout: float | None = None) -> int | None: + self.waited = True + return self.poll() + + +def _start(tmp_path, *, ready_at: float | None, exit_at: float | None, **budget): + """Run the real ``start`` against a service ready at ``ready_at``, then close. + + Returns whether the handle was available and what reached the service + before ``close``. + """ + + clock = _Clock() + process = _ServiceProcess(clock, exit_at) + sent: list[dict[str, object]] = [] + + def make_runtime_directory(**_kwargs) -> str: + runtime = tmp_path / f"runtime-{next(_RUNTIME_DIRECTORIES)}" + runtime.mkdir() + (runtime / SOCKET_FILENAME).touch() + return str(runtime) + + def service_ready(_socket_path) -> bool: + return process.poll() is None and ready_at is not None and clock.now >= ready_at + + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + module, + "time", + SimpleNamespace(monotonic=clock.monotonic, sleep=clock.sleep), + ) + patch.setattr(module.tempfile, "mkdtemp", make_runtime_directory) + patch.setattr(module.subprocess, "Popen", lambda *args, **kwargs: process) + patch.setattr(module, "_service_ready", service_ready) + patch.setattr(module, "runtime_identity", dict) + patch.setattr( + LocalTelemetryEmitter, "_send", lambda _self, payload: sent.append(payload) + ) + emitter = LocalTelemetryEmitter.start( + run_id="startup-wait", + country_code="US", + pipeline="test-pipeline", + development_collector_url="http://127.0.0.1:1", + spool_path=tmp_path / "events.sqlite3", + **budget, + ) + available = emitter.available + sent_before_close = list(sent) + emitter.close() + return available, process, clock, sent_before_close + + +_SECONDS = st.floats(min_value=0, max_value=10, allow_nan=False) + + +@settings( + max_examples=200, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +@given( + ready_at=st.none() | _SECONDS, + exit_at=st.none() | _SECONDS, + timeout=_SECONDS, +) +def test_start_keeps_exactly_the_services_that_answer_within_budget( + tmp_path, capsys, real_local_telemetry, ready_at, exit_at, timeout +) -> None: + events = [time for time in (ready_at, exit_at, timeout) if time is not None] + assume(all(abs(a - b) > _SEPARATION for a, b in itertools.combinations(events, 2))) + capsys.readouterr() + + available, process, clock, sent = _start( + tmp_path, ready_at=ready_at, exit_at=exit_at, startup_timeout_seconds=timeout + ) + + exits_first = exit_at is not None and exit_at < timeout + horizon = min(timeout, math.inf if exit_at is None else exit_at) + err = capsys.readouterr().err + if ready_at is not None and ready_at < horizon: + assert available + assert clock.now <= ready_at + STARTUP_POLL_SECONDS + 1e-9 + assert not process.terminated + assert sent[0]["event"]["message"] == BUILD_STARTED_MESSAGE + assert err == "" + else: + assert not available + assert clock.now <= horizon + STARTUP_POLL_SECONDS + 1e-9 + assert sent == [] + assert process.waited + if exits_first: + assert not process.terminated + assert f"exited with status {_EXIT_STATUS} before" in err + else: + assert process.terminated + assert f"did not become ready within {timeout:g} s" in err + + +def test_default_budget_keeps_a_service_that_starts_slowly( + tmp_path, capsys, real_local_telemetry +) -> None: + """A service needing 20 s, which the former 3 s budget abandoned, is kept.""" + + available, process, clock, _sent = _start(tmp_path, ready_at=20.0, exit_at=None) + + assert available + assert 20.0 <= clock.now <= 20.0 + STARTUP_POLL_SECONDS + 1e-9 + assert not process.terminated + assert capsys.readouterr().err == "" + + +def test_default_budget_abandons_a_service_that_never_answers( + tmp_path, capsys, real_local_telemetry +) -> None: + available, process, clock, sent = _start(tmp_path, ready_at=None, exit_at=None) + + assert not available + assert DEFAULT_STARTUP_TIMEOUT_SECONDS == 30.0 + assert clock.now <= DEFAULT_STARTUP_TIMEOUT_SECONDS + STARTUP_POLL_SECONDS + assert process.terminated and process.waited + assert sent == [] + assert "did not become ready within 30 s" in capsys.readouterr().err + + +def test_a_service_that_exits_is_reported_at_once( + tmp_path, capsys, real_local_telemetry +) -> None: + available, process, clock, sent = _start(tmp_path, ready_at=None, exit_at=0.5) + + assert not available + assert clock.now <= 0.5 + STARTUP_POLL_SECONDS + 1e-9 + assert not process.terminated + assert sent == [] + err = capsys.readouterr().err + assert f"exited with status {_EXIT_STATUS} before it became ready" in err + assert "within" not in err diff --git a/packages/microcosm-build/tests/engine_free/us/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/engine_free/us/test_us_modal_stage_plan_tool.py index 90307cd5fb..f6c5f6349e 100644 --- a/packages/microcosm-build/tests/engine_free/us/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/engine_free/us/test_us_modal_stage_plan_tool.py @@ -609,7 +609,7 @@ def test_runner_smoke_is_inline_and_check_sized() -> None: assert plan.resources is plan_lib.CHECK argv = plan_lib.planned_argv(plan) assert argv[:3] == ["/opt/venv/bin/python", "-B", "-c"] - assert "import microcosm.build" in argv[3] + assert "from microcosm.build import *" in argv[3] assert argv[4:] == [ "/work/state", "/work/inputs/ladder/us_puma_ladder_2020.npz", diff --git a/tools/modal_us_stage.py b/tools/modal_us_stage.py index 93d637e929..26f20f2a2e 100644 --- a/tools/modal_us_stage.py +++ b/tools/modal_us_stage.py @@ -646,7 +646,9 @@ def check_stage(plan_data: dict) -> dict: [ f"{plan_lib.IMAGE_VENV}/bin/python", "-c", - "import json, importlib.metadata as m; import microcosm.build; " + # The star import loads every public name; a bare package import + # is lazy and would not prove the modelling stack imports. + "import json, importlib.metadata as m; from microcosm.build import *; " "print(json.dumps({p: m.version(p) for p in " "('policyengine-us', 'policyengine-core', 'microcosm-build')}))", ], diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index 9a5cefeb23..7b0212468d 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -581,7 +581,9 @@ def _acs_local_release_argv( import pathlib import sys -import microcosm.build # noqa: F401 - proves the synced environment imports +# Every public name, not the lazy package alone: proves the synced +# environment's modelling stack (torch, pandas, calibration) imports. +from microcosm.build import * # noqa: F403 state = pathlib.Path(sys.argv[1]) rows = [