mostlyright-data 0.9.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- mostlyright/data_harness/__init__.py +158 -0
- mostlyright/data_harness/acquisition/__init__.py +55 -0
- mostlyright/data_harness/acquisition/http.py +2773 -0
- mostlyright/data_harness/acquisition/parsing.py +809 -0
- mostlyright/data_harness/acquisition/ranges.py +495 -0
- mostlyright/data_harness/acquisition/result_download.py +360 -0
- mostlyright/data_harness/acquisition/retention_admission.py +248 -0
- mostlyright/data_harness/acquisition/sandbox.py +4888 -0
- mostlyright/data_harness/acquisition/url_policy.py +530 -0
- mostlyright/data_harness/agent_runtime.py +2743 -0
- mostlyright/data_harness/assets/logo-ink.svg +31 -0
- mostlyright/data_harness/backends/__init__.py +28 -0
- mostlyright/data_harness/backends/pandas_backend.py +350 -0
- mostlyright/data_harness/backends/polars_backend.py +366 -0
- mostlyright/data_harness/backends/protocol.py +124 -0
- mostlyright/data_harness/backends/reference.py +83 -0
- mostlyright/data_harness/backends/registry.py +55 -0
- mostlyright/data_harness/backends/restrictions.py +126 -0
- mostlyright/data_harness/canonical.py +333 -0
- mostlyright/data_harness/catalog_job.py +625 -0
- mostlyright/data_harness/cli.py +5398 -0
- mostlyright/data_harness/contracts.py +53 -0
- mostlyright/data_harness/coordinator.py +1307 -0
- mostlyright/data_harness/deploy.py +924 -0
- mostlyright/data_harness/deploy_target.py +312 -0
- mostlyright/data_harness/deployment_evidence.py +1067 -0
- mostlyright/data_harness/event_presentation.py +576 -0
- mostlyright/data_harness/events.py +2152 -0
- mostlyright/data_harness/fast_delimited.py +239 -0
- mostlyright/data_harness/fleet.py +237 -0
- mostlyright/data_harness/formats.py +236 -0
- mostlyright/data_harness/governors.py +1163 -0
- mostlyright/data_harness/hosted_bootstrap.py +972 -0
- mostlyright/data_harness/hosted_crawler.py +1115 -0
- mostlyright/data_harness/hosted_crawler_container_smoke.py +351 -0
- mostlyright/data_harness/hosted_crawler_fetch.py +423 -0
- mostlyright/data_harness/hosted_crawler_job.py +1277 -0
- mostlyright/data_harness/hosted_crawler_protocol.py +676 -0
- mostlyright/data_harness/hosted_dataset.py +1500 -0
- mostlyright/data_harness/hosted_deploy.py +3037 -0
- mostlyright/data_harness/hosted_handoff.py +62 -0
- mostlyright/data_harness/hosted_ingestion_contract.py +504 -0
- mostlyright/data_harness/hosted_ingestion_job.py +356 -0
- mostlyright/data_harness/hosted_ingestion_job_smoke.py +40 -0
- mostlyright/data_harness/hosted_session_container_smoke.py +194 -0
- mostlyright/data_harness/hosted_session_worker.py +3554 -0
- mostlyright/data_harness/hosted_session_worker_job_smoke.py +46 -0
- mostlyright/data_harness/hosted_worker.py +6784 -0
- mostlyright/data_harness/ingestion/__init__.py +56 -0
- mostlyright/data_harness/ingestion/contracts.py +461 -0
- mostlyright/data_harness/ingestion/faults.py +42 -0
- mostlyright/data_harness/ingestion/gcs_store.py +1162 -0
- mostlyright/data_harness/ingestion/spool.py +130 -0
- mostlyright/data_harness/ingestion/store.py +885 -0
- mostlyright/data_harness/key_seam.py +434 -0
- mostlyright/data_harness/linux_process_boundary.py +262 -0
- mostlyright/data_harness/local_contracts.py +2880 -0
- mostlyright/data_harness/local_search/__init__.py +5 -0
- mostlyright/data_harness/local_search/build_index.py +1087 -0
- mostlyright/data_harness/local_search/contracts.py +920 -0
- mostlyright/data_harness/local_search/query_trace.py +266 -0
- mostlyright/data_harness/local_search/retrieval.py +700 -0
- mostlyright/data_harness/local_search/sealed.py +474 -0
- mostlyright/data_harness/local_search/service.py +784 -0
- mostlyright/data_harness/nbrender/CONTRACT.md +212 -0
- mostlyright/data_harness/nbrender/__init__.py +12 -0
- mostlyright/data_harness/nbrender/chrome.py +359 -0
- mostlyright/data_harness/nbrender/code_body.py +266 -0
- mostlyright/data_harness/nbrender/document.py +407 -0
- mostlyright/data_harness/nbrender/frame.py +275 -0
- mostlyright/data_harness/nbrender/interactive.py +337 -0
- mostlyright/data_harness/nbrender/markdown_body.py +477 -0
- mostlyright/data_harness/nbrender/mr_components.py +134 -0
- mostlyright/data_harness/nbrender/outputs_data.py +595 -0
- mostlyright/data_harness/nbrender/outputs_rich.py +906 -0
- mostlyright/data_harness/nbrender/outputs_source.py +260 -0
- mostlyright/data_harness/nbrender/outputs_stage.py +176 -0
- mostlyright/data_harness/nbrender/outputs_text.py +400 -0
- mostlyright/data_harness/nbrender/parse.py +394 -0
- mostlyright/data_harness/nbrender/status.py +40 -0
- mostlyright/data_harness/nbrender/tokens.py +1295 -0
- mostlyright/data_harness/notebook.py +1710 -0
- mostlyright/data_harness/offline.py +2049 -0
- mostlyright/data_harness/operation_registry.py +1007 -0
- mostlyright/data_harness/operator_setup.py +239 -0
- mostlyright/data_harness/pipeline.py +6428 -0
- mostlyright/data_harness/plan_graph.py +2026 -0
- mostlyright/data_harness/preparation/__init__.py +104 -0
- mostlyright/data_harness/preparation/contracts.py +1017 -0
- mostlyright/data_harness/preparation/engine.py +221 -0
- mostlyright/data_harness/preparation/errors.py +14 -0
- mostlyright/data_harness/preparation/gates.py +751 -0
- mostlyright/data_harness/preparation/joins.py +574 -0
- mostlyright/data_harness/preparation/profile.py +384 -0
- mostlyright/data_harness/preparation/table.py +217 -0
- mostlyright/data_harness/preparation/transforms.py +568 -0
- mostlyright/data_harness/progress_events.py +534 -0
- mostlyright/data_harness/readers/__init__.py +46 -0
- mostlyright/data_harness/readers/containers.py +963 -0
- mostlyright/data_harness/readers/contracts.py +542 -0
- mostlyright/data_harness/readers/delimited.py +257 -0
- mostlyright/data_harness/readers/grib2/__init__.py +33 -0
- mostlyright/data_harness/readers/grib2/admission.py +722 -0
- mostlyright/data_harness/readers/grib2/decode.py +1009 -0
- mostlyright/data_harness/readers/grib2/geometry.py +1133 -0
- mostlyright/data_harness/readers/grib2/portable_math.py +501 -0
- mostlyright/data_harness/readers/json_tabular.py +485 -0
- mostlyright/data_harness/readers/registry.py +514 -0
- mostlyright/data_harness/readers/samples/README.md +110 -0
- mostlyright/data_harness/readers/samples/archive.gzip/1.0.0/cities_one_stream/cities.csv.gz +0 -0
- mostlyright/data_harness/readers/samples/archive.gzip/1.0.0/cities_one_stream/expected.json +24 -0
- mostlyright/data_harness/readers/samples/archive.gzip/1.1.0/cities_one_stream/cities.csv.gz +0 -0
- mostlyright/data_harness/readers/samples/archive.gzip/1.1.0/cities_one_stream/expected.json +24 -0
- mostlyright/data_harness/readers/samples/archive.tar/1.0.0/cities_beside_a_directory_entry/cities.tar +0 -0
- mostlyright/data_harness/readers/samples/archive.tar/1.0.0/cities_beside_a_directory_entry/expected.json +24 -0
- mostlyright/data_harness/readers/samples/archive.tar/1.1.0/cities_beside_a_directory_entry/cities.tar +0 -0
- mostlyright/data_harness/readers/samples/archive.tar/1.1.0/cities_beside_a_directory_entry/expected.json +24 -0
- mostlyright/data_harness/readers/samples/archive.zip/1.0.0/cities_beside_a_second_member/cities.zip +0 -0
- mostlyright/data_harness/readers/samples/archive.zip/1.0.0/cities_beside_a_second_member/expected.json +25 -0
- mostlyright/data_harness/readers/samples/archive.zip/1.1.0/dwd_semicolon_station_member/dwd-station.zip +0 -0
- mostlyright/data_harness/readers/samples/archive.zip/1.1.0/dwd_semicolon_station_member/expected.json +25 -0
- mostlyright/data_harness/readers/samples/archive.zip/1.2.0/dwd_semicolon_station_member/dwd-station.zip +0 -0
- mostlyright/data_harness/readers/samples/archive.zip/1.2.0/dwd_semicolon_station_member/expected.json +25 -0
- mostlyright/data_harness/readers/samples/delimited_text/1.0.0/an_ordinary_comma_separated_table/cities.csv +3 -0
- mostlyright/data_harness/readers/samples/delimited_text/1.0.0/an_ordinary_comma_separated_table/expected.json +23 -0
- mostlyright/data_harness/readers/samples/delimited_text/1.0.0/quoted_fields_holding_the_delimiter/cities.tsv +5 -0
- mostlyright/data_harness/readers/samples/delimited_text/1.0.0/quoted_fields_holding_the_delimiter/expected.json +25 -0
- mostlyright/data_harness/readers/samples/delimited_text/1.1.0/an_hourly_observation_table_served_as_plain_text/expected.json +30 -0
- mostlyright/data_harness/readers/samples/delimited_text/1.1.0/an_hourly_observation_table_served_as_plain_text/observations.csv +5 -0
- mostlyright/data_harness/readers/samples/json.tabular/1.0.0/nested_hourly_observations/expected.json +44 -0
- mostlyright/data_harness/readers/samples/json.tabular/1.0.0/nested_hourly_observations/stations.json +1 -0
- mostlyright/data_harness/readers/samples/json.tabular/1.1.0/an_observation_stream_served_as_plain_text/expected.json +48 -0
- mostlyright/data_harness/readers/samples/json.tabular/1.1.0/an_observation_stream_served_as_plain_text/observations.ndjson +4 -0
- mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.0.0/an_ordinary_table_beside_a_second_sheet/cities.xlsx +0 -0
- mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.0.0/an_ordinary_table_beside_a_second_sheet/expected.json +24 -0
- mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.0.0/shares_the_workbook_had_already_computed/expected.json +27 -0
- mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.0.0/shares_the_workbook_had_already_computed/shares.xlsx +0 -0
- mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.1.0/shares_the_workbook_had_already_computed/expected.json +27 -0
- mostlyright/data_harness/readers/samples/spreadsheet.xlsx/1.1.0/shares_the_workbook_had_already_computed/shares.xlsx +0 -0
- mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/README.md +20 -0
- mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/gfs_2m_temperature/expected.json +55 -0
- mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/gfs_2m_temperature/gfs-2m-temperature.grib2 +0 -0
- mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/hrrr_2m_temperature/expected.json +54 -0
- mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/hrrr_2m_temperature/hrrr-2m-temperature.grib2 +0 -0
- mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/hrrr_categorical_rain/expected.json +54 -0
- mostlyright/data_harness/readers/samples/weather.grib2/1.0.0/hrrr_categorical_rain/hrrr-categorical-rain.grib2 +0 -0
- mostlyright/data_harness/readers/samples/weather.grib2/2.0.0/hrrr_2m_temperature/expected.json +54 -0
- mostlyright/data_harness/readers/samples/weather.grib2/2.0.0/hrrr_2m_temperature/hrrr-2m-temperature.grib2 +0 -0
- mostlyright/data_harness/readers/samples.py +582 -0
- mostlyright/data_harness/readers/spreadsheet.py +803 -0
- mostlyright/data_harness/readers/tabular.py +510 -0
- mostlyright/data_harness/recipe.py +5321 -0
- mostlyright/data_harness/repair/__init__.py +78 -0
- mostlyright/data_harness/repair/adapters.py +274 -0
- mostlyright/data_harness/repair/contracts.py +872 -0
- mostlyright/data_harness/repair/coordinator.py +1099 -0
- mostlyright/data_harness/repair/errors.py +16 -0
- mostlyright/data_harness/review.py +2533 -0
- mostlyright/data_harness/rowset.py +283 -0
- mostlyright/data_harness/serving.py +1975 -0
- mostlyright/data_harness/serving_edge.py +590 -0
- mostlyright/data_harness/serving_http.py +1031 -0
- mostlyright/data_harness/session_probes.py +759 -0
- mostlyright/data_harness/signing.py +101 -0
- mostlyright/data_harness/source_discovery.py +898 -0
- mostlyright/data_harness/sources/__init__.py +209 -0
- mostlyright/data_harness/sources/_adapter_steps.py +213 -0
- mostlyright/data_harness/sources/adapters.py +1214 -0
- mostlyright/data_harness/sources/cadence.py +1428 -0
- mostlyright/data_harness/sources/cadence_emission.py +453 -0
- mostlyright/data_harness/sources/cadence_history.py +546 -0
- mostlyright/data_harness/sources/catalog/__init__.py +17 -0
- mostlyright/data_harness/sources/catalog/admission.py +477 -0
- mostlyright/data_harness/sources/catalog/authoring.py +1701 -0
- mostlyright/data_harness/sources/catalog/authoring_policy.py +701 -0
- mostlyright/data_harness/sources/catalog/authoring_shards.py +1217 -0
- mostlyright/data_harness/sources/catalog/bounded_io.py +231 -0
- mostlyright/data_harness/sources/catalog/channel.py +523 -0
- mostlyright/data_harness/sources/catalog/channel_client.py +296 -0
- mostlyright/data_harness/sources/catalog/contracts.py +825 -0
- mostlyright/data_harness/sources/catalog/coverage.py +137 -0
- mostlyright/data_harness/sources/catalog/delta.py +1340 -0
- mostlyright/data_harness/sources/catalog/embedding.py +532 -0
- mostlyright/data_harness/sources/catalog/entry_v2.py +1182 -0
- mostlyright/data_harness/sources/catalog/fill.py +3889 -0
- mostlyright/data_harness/sources/catalog/fill_partitions.py +459 -0
- mostlyright/data_harness/sources/catalog/fill_staging.py +1105 -0
- mostlyright/data_harness/sources/catalog/gating.py +374 -0
- mostlyright/data_harness/sources/catalog/generation_receipt.py +1607 -0
- mostlyright/data_harness/sources/catalog/harvest/__init__.py +7 -0
- mostlyright/data_harness/sources/catalog/harvest/ckan.py +384 -0
- mostlyright/data_harness/sources/catalog/harvest/datagov_v4.py +798 -0
- mostlyright/data_harness/sources/catalog/harvest/protocol.py +964 -0
- mostlyright/data_harness/sources/catalog/harvest/sdmx.py +445 -0
- mostlyright/data_harness/sources/catalog/harvest/stac.py +384 -0
- mostlyright/data_harness/sources/catalog/health.py +447 -0
- mostlyright/data_harness/sources/catalog/hosted_catalog.py +105 -0
- mostlyright/data_harness/sources/catalog/identity_history.py +1549 -0
- mostlyright/data_harness/sources/catalog/neural.py +1618 -0
- mostlyright/data_harness/sources/catalog/packed_catalog.py +2345 -0
- mostlyright/data_harness/sources/catalog/packed_retrieval.py +1517 -0
- mostlyright/data_harness/sources/catalog/packed_writer.py +2802 -0
- mostlyright/data_harness/sources/catalog/query_trace.py +1037 -0
- mostlyright/data_harness/sources/catalog/recommend.py +171 -0
- mostlyright/data_harness/sources/catalog/retrieval.py +230 -0
- mostlyright/data_harness/sources/catalog/retrieval_manifest.py +995 -0
- mostlyright/data_harness/sources/catalog/rights_decisions.py +254 -0
- mostlyright/data_harness/sources/catalog/sealed.py +560 -0
- mostlyright/data_harness/sources/catalog/search.py +230 -0
- mostlyright/data_harness/sources/catalog/streaming_delta.py +1097 -0
- mostlyright/data_harness/sources/catalog/update.py +891 -0
- mostlyright/data_harness/sources/collections.py +815 -0
- mostlyright/data_harness/sources/contracts.py +2223 -0
- mostlyright/data_harness/sources/deletion.py +761 -0
- mostlyright/data_harness/sources/fitness.py +162 -0
- mostlyright/data_harness/sources/governance.py +163 -0
- mostlyright/data_harness/sources/hosted.py +173 -0
- mostlyright/data_harness/sources/integration.py +218 -0
- mostlyright/data_harness/sources/range_reader.py +418 -0
- mostlyright/data_harness/sources/registry.py +514 -0
- mostlyright/data_harness/sources/rights_rule.py +59 -0
- mostlyright/data_harness/sources/source_cadence_vectors.v1.json +1 -0
- mostlyright/data_harness/sources/sports.py +521 -0
- mostlyright/data_harness/sources/stream.py +524 -0
- mostlyright/data_harness/sources/stream_connector.py +418 -0
- mostlyright/data_harness/sources/stream_recorder.py +1404 -0
- mostlyright/data_harness/studio_boundary.py +2019 -0
- mostlyright/data_harness/thin/__init__.py +37 -0
- mostlyright/data_harness/thin/acquire.py +1137 -0
- mostlyright/data_harness/thin/acquire_cancel.py +579 -0
- mostlyright/data_harness/thin/approvals.py +617 -0
- mostlyright/data_harness/thin/commands.py +406 -0
- mostlyright/data_harness/thin/download.py +194 -0
- mostlyright/data_harness/thin/narrative.py +589 -0
- mostlyright/data_harness/thin/parity.py +1070 -0
- mostlyright/data_harness/thin/propose.py +2759 -0
- mostlyright/data_harness/thin/research.py +1663 -0
- mostlyright/data_harness/thin/router.py +924 -0
- mostlyright/data_harness/thin/runs.py +519 -0
- mostlyright/data_harness/thin/session.py +281 -0
- mostlyright/data_harness/thin/stream.py +501 -0
- mostlyright/data_harness/thin/transport.py +187 -0
- mostlyright/data_harness/thin/vocabulary.py +368 -0
- mostlyright/data_harness/thin/workers.py +164 -0
- mostlyright/data_harness/ucum/TABLE-PIN.json +40 -0
- mostlyright/data_harness/ucum/ucum-subset.v1.json +632 -0
- mostlyright/data_harness/unit_flow.py +927 -0
- mostlyright/data_harness/units.py +572 -0
- mostlyright/data_harness/ux/__init__.py +9 -0
- mostlyright/data_harness/ux/approve.py +485 -0
- mostlyright/data_harness/ux/author_yaml.py +597 -0
- mostlyright/data_harness/ux/cloud_auth.py +447 -0
- mostlyright/data_harness/ux/commands/__init__.py +260 -0
- mostlyright/data_harness/ux/commands/approve.py +136 -0
- mostlyright/data_harness/ux/commands/auth.py +744 -0
- mostlyright/data_harness/ux/commands/author.py +79 -0
- mostlyright/data_harness/ux/commands/catalog_author.py +403 -0
- mostlyright/data_harness/ux/commands/catalog_fill.py +523 -0
- mostlyright/data_harness/ux/commands/catalog_harvest.py +545 -0
- mostlyright/data_harness/ux/commands/catalog_publish.py +1838 -0
- mostlyright/data_harness/ux/commands/catalog_search.py +71 -0
- mostlyright/data_harness/ux/commands/catalog_update.py +437 -0
- mostlyright/data_harness/ux/commands/deploy.py +134 -0
- mostlyright/data_harness/ux/commands/deploy_dataset.py +98 -0
- mostlyright/data_harness/ux/commands/deploy_plan.py +105 -0
- mostlyright/data_harness/ux/commands/deploy_status.py +104 -0
- mostlyright/data_harness/ux/commands/diff.py +74 -0
- mostlyright/data_harness/ux/commands/index.py +84 -0
- mostlyright/data_harness/ux/commands/inventory.py +47 -0
- mostlyright/data_harness/ux/commands/list_builds.py +143 -0
- mostlyright/data_harness/ux/commands/login.py +63 -0
- mostlyright/data_harness/ux/commands/peek.py +236 -0
- mostlyright/data_harness/ux/commands/plan_check.py +90 -0
- mostlyright/data_harness/ux/commands/preflight.py +97 -0
- mostlyright/data_harness/ux/commands/record.py +107 -0
- mostlyright/data_harness/ux/commands/review_setup.py +47 -0
- mostlyright/data_harness/ux/commands/search.py +440 -0
- mostlyright/data_harness/ux/commands/show.py +61 -0
- mostlyright/data_harness/ux/commands/whoami.py +37 -0
- mostlyright/data_harness/ux/credential_native.py +551 -0
- mostlyright/data_harness/ux/credential_store.py +1055 -0
- mostlyright/data_harness/ux/credentials.py +631 -0
- mostlyright/data_harness/ux/diffing.py +444 -0
- mostlyright/data_harness/ux/headline.py +671 -0
- mostlyright/data_harness/ux/hosted_acquisition.py +974 -0
- mostlyright/data_harness/ux/hosted_run_status.py +619 -0
- mostlyright/data_harness/ux/inventory.py +427 -0
- mostlyright/data_harness/ux/local_review.py +375 -0
- mostlyright/data_harness/ux/login.py +691 -0
- mostlyright/data_harness/ux/path_kind.py +147 -0
- mostlyright/data_harness/ux/peek.py +1000 -0
- mostlyright/data_harness/ux/plain_file.py +178 -0
- mostlyright/data_harness/ux/plan_check.py +311 -0
- mostlyright/data_harness/ux/preflight.py +918 -0
- mostlyright/data_harness/ux/readers.py +1124 -0
- mostlyright/data_harness/ux/remediation.py +2195 -0
- mostlyright/data_harness/ux/render.py +657 -0
- mostlyright/data_harness/ux/workload.py +1077 -0
- mostlyright/data_harness/viewer.py +3713 -0
- mostlyright/data_harness/visual_run/__init__.py +83 -0
- mostlyright/data_harness/visual_run/authoring.py +235 -0
- mostlyright/data_harness/visual_run/contracts.py +673 -0
- mostlyright/data_harness/visual_run/materialize.py +486 -0
- mostlyright/data_harness/visual_run/observations.py +874 -0
- mostlyright/data_harness/visual_run/query.py +259 -0
- mostlyright/data_harness/visual_run/reducer.py +280 -0
- mostlyright/data_harness/visual_run/sdk.py +892 -0
- mostlyright/data_harness/visual_run/store.py +584 -0
- mostlyright/data_harness/visual_run/transport.py +239 -0
- mostlyright/data_harness/watch.py +2999 -0
- mostlyright_data-0.9.0.dist-info/METADATA +607 -0
- mostlyright_data-0.9.0.dist-info/RECORD +314 -0
- mostlyright_data-0.9.0.dist-info/WHEEL +4 -0
- mostlyright_data-0.9.0.dist-info/entry_points.txt +12 -0
|
@@ -0,0 +1,3554 @@
|
|
|
1
|
+
"""The warm worker behind one research session (ADR 0021).
|
|
2
|
+
|
|
3
|
+
A session worker is not a new kind of worker. Studio's ``RESEARCH-SESSIONS.md`` is explicit about
|
|
4
|
+
it: *"A session worker IS a producer-role attempt on the session run"*, with ``worker_kind =
|
|
5
|
+
"session"`` and two extra capability scopes, so everything it does after the lease --
|
|
6
|
+
``attempt:start``, batch ``event:append``, ``attempt:heartbeat``, ``attempt:fail`` -- is an existing
|
|
7
|
+
fenced producer callback with the existing workspace/run/attempt/generation/fence binding. Only
|
|
8
|
+
four verbs are new -- ``probes:claim``, ``probes/{id}:result``, and the ``session-acquisitions``
|
|
9
|
+
pair that gives a probe its bytes -- and this module is the whole of the harness half of that
|
|
10
|
+
protocol.
|
|
11
|
+
|
|
12
|
+
Read this module against :mod:`mostlyright.data_harness.hosted_bootstrap`, which establishes the
|
|
13
|
+
same workload-identity and bounded-secret handling discipline.
|
|
14
|
+
|
|
15
|
+
**What is kept.** Ambient workload identity is read once, from the metadata server, into a
|
|
16
|
+
``bytearray`` that is zeroed in a ``finally``. The capability that comes back is never a ``str``
|
|
17
|
+
this module holds: it goes straight into a :class:`~mostlyright.data_harness.hosted_worker.
|
|
18
|
+
CredentialLease`, which is non-serializable, redacted in every repr, and zeroed on close. There is
|
|
19
|
+
exactly one capability per attempt and exactly one attempt per session, so "one capability per
|
|
20
|
+
attempt" survives verbatim. Nothing is written to disk: the lease response is parsed out of a
|
|
21
|
+
bounded in-memory buffer that is zeroed as soon as the capability is inside its lease.
|
|
22
|
+
|
|
23
|
+
**What is different, and why it has to be.** The generic bootstrap reads a dispatch descriptor from
|
|
24
|
+
a file descriptor because producer and verifier jobs carry sealed job descriptions. A research
|
|
25
|
+
session Job instead receives three non-secret arguments and uses the atomically claimed dispatch
|
|
26
|
+
record in Studio as its bootstrap authority. A capability-free monitor starts an isolated child;
|
|
27
|
+
only that child exchanges and retains the capability in ``CredentialLease`` for the session
|
|
28
|
+
lifetime, never in argv, environment or a descriptor visible to another task.
|
|
29
|
+
|
|
30
|
+
**Request lifetime, and how scale-to-zero actually happens.** This is a Cloud Run *Job*, not a
|
|
31
|
+
service. Studio supplies only an opaque UUID4 dispatch id, an exact Studio exchange URL and the OIDC
|
|
32
|
+
audience as task arguments. The task reads its ambient workload identity, atomically exchanges the
|
|
33
|
+
dispatch id for the one fenced session lease, and runs :class:`SessionSupervisor` synchronously to
|
|
34
|
+
its terminal state. The Job task is therefore the unit Cloud Run retains: an overlapping session
|
|
35
|
+
gets a distinct task and cannot share a warm service instance or an in-memory capability with the
|
|
36
|
+
first.
|
|
37
|
+
|
|
38
|
+
The dispatch id is deliberately not a lease nonce. It is non-secret routing data, accepted exactly
|
|
39
|
+
once by Studio's transaction that creates and fences the attempt. The capability appears only in the
|
|
40
|
+
exchange response, immediately enters ``CredentialLease``, and is zeroed when the supervisor exits.
|
|
41
|
+
Job retry policy and task timeout are deployment controls, not a second leasing protocol: a retry
|
|
42
|
+
of an already-claimed dispatch receives Studio's typed duplicate/terminal refusal and must not run a
|
|
43
|
+
second supervisor.
|
|
44
|
+
|
|
45
|
+
**Exit is by stopping, never by completing.** There is no ``:complete`` here. Studio's contract
|
|
46
|
+
says why: ``:complete`` requires an admitted candidate, which a session never has, and after a
|
|
47
|
+
close the session run is already terminal so ``:fail`` cannot land either. ``:fail`` is reserved
|
|
48
|
+
for the one case where it is still meaningful -- the worker dying while its session is still live.
|
|
49
|
+
|
|
50
|
+
**Heartbeat is liveness, never activity.** It runs on its own timer thread and it deliberately does
|
|
51
|
+
not touch ``last_activity_at`` server-side. That is Studio's rule and the reason is the bill: a
|
|
52
|
+
worker heartbeats every thirty seconds, so counting a heartbeat as activity would mean no session
|
|
53
|
+
ever goes idle and scale-to-session never happens. Nothing in this module tries to keep a session
|
|
54
|
+
alive; keeping a session alive is the user's job, by using it.
|
|
55
|
+
"""
|
|
56
|
+
|
|
57
|
+
from __future__ import annotations
|
|
58
|
+
|
|
59
|
+
import contextlib
|
|
60
|
+
import hashlib
|
|
61
|
+
import json
|
|
62
|
+
import logging
|
|
63
|
+
import math
|
|
64
|
+
import os
|
|
65
|
+
import re
|
|
66
|
+
import select
|
|
67
|
+
import signal
|
|
68
|
+
import subprocess
|
|
69
|
+
import sys
|
|
70
|
+
import tempfile
|
|
71
|
+
import threading
|
|
72
|
+
import time
|
|
73
|
+
from collections.abc import Callable, Mapping, Sequence
|
|
74
|
+
from dataclasses import dataclass, field, replace
|
|
75
|
+
from datetime import UTC, datetime
|
|
76
|
+
from http import HTTPStatus
|
|
77
|
+
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
78
|
+
from pathlib import Path
|
|
79
|
+
from typing import Any, ClassVar, Protocol
|
|
80
|
+
from uuid import UUID
|
|
81
|
+
|
|
82
|
+
from mostlyright.data_harness import progress_events, session_probes
|
|
83
|
+
from mostlyright.data_harness.acquisition.parsing import ParseLimits
|
|
84
|
+
from mostlyright.data_harness.acquisition.result_download import (
|
|
85
|
+
ResultDownloadError,
|
|
86
|
+
validate_result_download,
|
|
87
|
+
verify_result_bytes,
|
|
88
|
+
)
|
|
89
|
+
from mostlyright.data_harness.canonical import canonical_json_bytes, sha256_bytes
|
|
90
|
+
from mostlyright.data_harness.hosted_bootstrap import (
|
|
91
|
+
BootstrapTransport,
|
|
92
|
+
UrlLibBootstrapTransport,
|
|
93
|
+
_https_url,
|
|
94
|
+
_metadata_identity,
|
|
95
|
+
)
|
|
96
|
+
from mostlyright.data_harness.hosted_crawler_protocol import (
|
|
97
|
+
HostedCrawlerProtocolError,
|
|
98
|
+
parse_crawler_result,
|
|
99
|
+
)
|
|
100
|
+
from mostlyright.data_harness.hosted_worker import (
|
|
101
|
+
MAX_CAPABILITY_BYTES,
|
|
102
|
+
CredentialLease,
|
|
103
|
+
HostedWorkerError,
|
|
104
|
+
_authenticated_client,
|
|
105
|
+
_idempotency,
|
|
106
|
+
_load_generated,
|
|
107
|
+
_new_transfer_client,
|
|
108
|
+
_ProducerAuthority,
|
|
109
|
+
_response,
|
|
110
|
+
_strict_json,
|
|
111
|
+
parse_capability,
|
|
112
|
+
)
|
|
113
|
+
from mostlyright.data_harness.linux_process_boundary import (
|
|
114
|
+
_child_subreaper_enabled,
|
|
115
|
+
_set_child_subreaper,
|
|
116
|
+
)
|
|
117
|
+
|
|
118
|
+
LOGGER = logging.getLogger(__name__)
|
|
119
|
+
|
|
120
|
+
#: The role this worker presents. Studio mints the attempt with ``worker_kind = "session"``; the
|
|
121
|
+
#: value is restated here only for the progress record that names the role.
|
|
122
|
+
SESSION_WORKER_ROLE = "session"
|
|
123
|
+
|
|
124
|
+
#: The one path Studio invokes, spelled by ``GoogleCloudRunSessionWorkerStarter``.
|
|
125
|
+
START_PATH = "/sessions:start"
|
|
126
|
+
#: An internal, descriptor-only command used by the retained Service rollback path. The service
|
|
127
|
+
#: process starts a fresh interpreter rather than forking Python from a request thread; the only
|
|
128
|
+
#: value passed in argv is an inherited descriptor number, never a nonce or capability.
|
|
129
|
+
LEGACY_BOOTSTRAP_FD_ARGUMENT = "--legacy-bootstrap-fd"
|
|
130
|
+
#: Liveness only. Cloud Run's startup probe needs an endpoint that answers before a session exists,
|
|
131
|
+
#: and answering it must never imply a session is warm -- ``warming`` is Studio's state to clear,
|
|
132
|
+
#: on a fenced attempt callback, and a 200 from here is not one.
|
|
133
|
+
HEALTH_PATHS = ("/livez", "/healthz")
|
|
134
|
+
|
|
135
|
+
#: Studio's start invocation carries its own contract version.
|
|
136
|
+
START_SCHEMA_VERSION = "3.0.0"
|
|
137
|
+
#: Cloud Run Job dispatches use the same schema as the lease response they mint. The dispatch
|
|
138
|
+
#: itself is deliberately only an opaque routing token; it is never a capability or lease nonce.
|
|
139
|
+
SESSION_DISPATCH_SCHEMA_VERSION = "3.0.0"
|
|
140
|
+
#: The only Studio URL a Job task may exchange its dispatch id against. The id is part of the
|
|
141
|
+
#: path so a task cannot turn an opaque value into a request to a neighbouring session.
|
|
142
|
+
SESSION_DISPATCH_EXCHANGE_PATH_TEMPLATE = (
|
|
143
|
+
"/internal/v3/session-worker/dispatches/{dispatch_id}:exchange"
|
|
144
|
+
)
|
|
145
|
+
#: Studio refuses an exchange after its one dispatch has already been consumed. These are normal
|
|
146
|
+
#: no-op terminal outcomes for a Cloud Run Job task, not failed executions: the task must not
|
|
147
|
+
#: invoke ``attempt:start`` after another task has claimed it or after Studio has closed it.
|
|
148
|
+
CONSUMED_SESSION_DISPATCH_CODES = frozenset(
|
|
149
|
+
{
|
|
150
|
+
"RESEARCH_SESSION_DISPATCH_ALREADY_CLAIMED",
|
|
151
|
+
"RESEARCH_SESSION_DISPATCH_CLOSED",
|
|
152
|
+
}
|
|
153
|
+
)
|
|
154
|
+
#: The lease path Studio publishes. Derived here rather than taken from the invocation; see
|
|
155
|
+
#: :func:`parse_start_invocation`.
|
|
156
|
+
LEASE_PATH_TEMPLATE = "/internal/v3/session-worker/sessions/{session_id}:lease"
|
|
157
|
+
|
|
158
|
+
#: A start invocation is five short fields. Anything larger is not one.
|
|
159
|
+
MAX_START_REQUEST_BYTES = 8 * 1024
|
|
160
|
+
#: A lease has no source bytes and must remain a small control-plane document. This separate
|
|
161
|
+
#: bound means a compromised exchange cannot make a long-lived Job hold an arbitrary response.
|
|
162
|
+
MAX_SESSION_DISPATCH_RESPONSE_BYTES = 64 * 1024
|
|
163
|
+
#: The service-to-fresh-interpreter start document is an already validated start request and must
|
|
164
|
+
#: stay no larger than the public request surface. The child accepts exactly one such document.
|
|
165
|
+
MAX_LEGACY_BOOTSTRAP_BYTES = MAX_START_REQUEST_BYTES
|
|
166
|
+
#: Service start must not acknowledge until its fresh monitor has installed the cancellation
|
|
167
|
+
#: handler that owns an inner capability child. This remains far below Studio's start request
|
|
168
|
+
#: timeout while rejecting a stuck/interrupted bootstrap rather than accepting a zombie session.
|
|
169
|
+
LEGACY_BOOTSTRAP_READY_TIMEOUT_SECONDS = 5.0
|
|
170
|
+
#: Studio bounds the nonce at 32..128 URL-safe characters.
|
|
171
|
+
MIN_NONCE_CHARS = 32
|
|
172
|
+
MAX_NONCE_CHARS = 128
|
|
173
|
+
|
|
174
|
+
#: How long the process waits for another session before exiting so Cloud Run can drop the
|
|
175
|
+
#: instance. Long enough that reopening a session is warm, short enough that a finished session is
|
|
176
|
+
#: not billed for a coffee break.
|
|
177
|
+
DEFAULT_SHUTDOWN_GRACE_SECONDS = 30.0
|
|
178
|
+
#: Pre-spawned no-network sandbox children. This is the first production caller of the warm pool
|
|
179
|
+
#: (#306), and the pool exists precisely for this: the follow-up-probe budget is 100-300 ms and a
|
|
180
|
+
#: cold interpreter spawn does not fit in it. On Linux under a bounded cgroup root
|
|
181
|
+
#: ``start_warm_pool`` returns 0 by design -- the count is read, not assumed, and a zero is not an
|
|
182
|
+
#: error, only a slower first probe.
|
|
183
|
+
DEFAULT_WARM_POOL_SIZE = 2
|
|
184
|
+
#: A probe is interactive. A probe that is not finished in this long is failed, not waited on.
|
|
185
|
+
DEFAULT_PROBE_TIMEOUT_SECONDS = 45.0
|
|
186
|
+
#: How long one session acquisition may take from the create call to a settled status. This is the
|
|
187
|
+
#: network half of a probe and it is bounded SEPARATELY from the Clean room ceiling above, because
|
|
188
|
+
#: the two fail for unrelated reasons and a client that is told "the crawl did not come back" can
|
|
189
|
+
#: act on that differently from "parsing did not finish".
|
|
190
|
+
DEFAULT_ACQUISITION_TIMEOUT_SECONDS = 90.0
|
|
191
|
+
#: How often the worker asks Studio whether its acquisition has settled. Fixed rather than backed
|
|
192
|
+
#: off: the whole wait is bounded above by one interactive ceiling, so a backoff would only make
|
|
193
|
+
#: the common fast case slower without making the slow case cheaper.
|
|
194
|
+
ACQUISITION_POLL_SECONDS = 1.0
|
|
195
|
+
#: The largest quarantine bundle a probe will pull into memory. The contract's own ceiling is 6 GiB
|
|
196
|
+
#: -- correct for a Build, absurd for a preview -- and the bundle is base64 inside canonical JSON,
|
|
197
|
+
#: so this is comfortably above `ProbeExecutor._max_source_bytes` (16 MiB) times the 4/3 encoding
|
|
198
|
+
#: expansion plus the receipt. A bundle over it is refused on its DECLARED size, before a byte is
|
|
199
|
+
#: transferred, which is the only place refusing is free.
|
|
200
|
+
MAX_PROBE_ACQUISITION_BUNDLE_BYTES = 32 * 1024 * 1024
|
|
201
|
+
#: Poll and heartbeat floors, used only when a lease omits them (it cannot -- both are required --
|
|
202
|
+
#: but a default beats a KeyError if Studio ever relaxes that).
|
|
203
|
+
FALLBACK_CLAIM_POLL_SECONDS = 1.0
|
|
204
|
+
FALLBACK_HEARTBEAT_SECONDS = 30.0
|
|
205
|
+
#: A refused claim poll is transient far more often than it is fatal. After this many in a row the
|
|
206
|
+
#: worker gives up on the session rather than polling a Studio that is saying no forever.
|
|
207
|
+
MAX_CONSECUTIVE_CLAIM_FAILURES = 10
|
|
208
|
+
#: Studio only lets a session shorten the ordinary thirty-minute producer lease. A later value
|
|
209
|
+
#: would make a compromised or mismatched control plane turn one session capability into an
|
|
210
|
+
#: unbounded worker lifetime, so it is a lease refusal rather than a value to clamp.
|
|
211
|
+
MAX_SESSION_CAPABILITY_REMAINING_SECONDS = 30 * 60
|
|
212
|
+
#: The isolated session child communicates only an expiry, identity and terminal summary to its
|
|
213
|
+
#: capability-free parent. It has no reason to emit a large record, and this ceiling makes the
|
|
214
|
+
#: parent-side parser a fixed-memory boundary if a child is compromised.
|
|
215
|
+
MAX_SESSION_SUPERVISOR_MESSAGE_BYTES = 128 * 1024
|
|
216
|
+
|
|
217
|
+
_UTC_TIMESTAMP = re.compile(
|
|
218
|
+
r"^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(?:\.[0-9]{1,6})?Z$"
|
|
219
|
+
)
|
|
220
|
+
|
|
221
|
+
#: The failure codes that mean the probe was REFUSED rather than attempted-and-broken. The
|
|
222
|
+
#: distinction is the only part of a failure a client can act on: a refused probe refuses the same
|
|
223
|
+
#: way every time, so the fix is a different probe -- or, for the source channel, somebody wiring
|
|
224
|
+
#: one -- while a broken probe may well not break twice. The set is written out rather than
|
|
225
|
+
#: inferred from the exception class, because the same class carries both halves.
|
|
226
|
+
PROBE_REFUSAL_CODES = frozenset(
|
|
227
|
+
{
|
|
228
|
+
# The open gap: no session-scoped source channel exists, so nothing was read.
|
|
229
|
+
"PROBE_SOURCE_CHANNEL_UNAVAILABLE",
|
|
230
|
+
# The session's own authority check, made before any read.
|
|
231
|
+
"PROBE_SOURCE_OUT_OF_SCOPE",
|
|
232
|
+
# A reader that answered for a different source. Refused before the bytes are used.
|
|
233
|
+
"PROBE_SOURCE_MISMATCH",
|
|
234
|
+
# The request body is not in the probe vocabulary.
|
|
235
|
+
"PROBE_REQUEST_INVALID",
|
|
236
|
+
# The request names a column, or an aggregate over a column, this source cannot answer.
|
|
237
|
+
# Discovered after the read rather than before it, and still a refusal of the REQUEST:
|
|
238
|
+
# re-asking it changes nothing.
|
|
239
|
+
"PROBE_COLUMN_UNKNOWN",
|
|
240
|
+
"PROBE_AGGREGATE_NOT_NUMERIC",
|
|
241
|
+
# Studio's own refusals on the session-acquisition door. The first is the session's
|
|
242
|
+
# source allowlist, checked before any lookup so a refusal is not an existence oracle;
|
|
243
|
+
# the second is "this source has no single active public.https registration", which is a
|
|
244
|
+
# fact about the source's registration and not about this probe.
|
|
245
|
+
"RESEARCH_SOURCE_FORBIDDEN",
|
|
246
|
+
"RESEARCH_SOURCE_UNACQUIRABLE",
|
|
247
|
+
# The acquired source is larger than a probe may hold. A property of the source.
|
|
248
|
+
"PROBE_SOURCE_TOO_LARGE",
|
|
249
|
+
}
|
|
250
|
+
)
|
|
251
|
+
#: The Clean room's own code for "did not finish in time" (`acquisition.sandbox`). It is distinct
|
|
252
|
+
#: from a source acquisition timing out: reducing a scan can help after bytes arrived, but cannot
|
|
253
|
+
#: make a delayed remote acquisition arrive sooner.
|
|
254
|
+
PROBE_TIMEOUT_CODES = frozenset({"SANDBOX_TIMEOUT", "PROBE_SOURCE_ACQUISITION_TIMEOUT"})
|
|
255
|
+
#: Studio's code for an acquisition whose own expiry passed -- the earliest of the crawler-session
|
|
256
|
+
#: TTL, the deployment's admission window and the worker's attempt lease. It is the one failure the
|
|
257
|
+
#: reader retries, because Studio releases the admission slot as it expires the acquisition, so
|
|
258
|
+
#: opening a fresh one is not a second spend. See :meth:`StudioSessionAcquisitionReader.read`.
|
|
259
|
+
ACQUISITION_EXPIRED_CODE = "CRAWLER_SESSION_EXPIRED"
|
|
260
|
+
#: A ceiling somebody else is holding, not a refusal of this request. Studio charges a session
|
|
261
|
+
#: acquisition to the principal who opened the session, against the same active-and-rate admission
|
|
262
|
+
#: limits a person acquiring by hand would spend (ADR 0015: two at once, four in sixty seconds), so
|
|
263
|
+
#: this arrives when the session is busy rather than when the probe is wrong. It is its own outcome
|
|
264
|
+
#: because the advice it carries is the opposite of a refusal's: ask again in a moment.
|
|
265
|
+
SESSION_ACQUISITION_THROTTLED = "SESSION_ACQUISITION_THROTTLED"
|
|
266
|
+
#: Studio's wire code is retained in the probe result. The progress/UI code is separate so an
|
|
267
|
+
#: acquisition admission window is not presented as model-token, billing, sign-in, or source-key
|
|
268
|
+
#: configuration trouble.
|
|
269
|
+
PROBE_THROTTLE_WIRE_CODES = frozenset({"AGENT_BUDGET_EXCEEDED"})
|
|
270
|
+
PROBE_THROTTLE_CODES = frozenset({SESSION_ACQUISITION_THROTTLED})
|
|
271
|
+
MAX_RETRY_AFTER_SECONDS = 3600
|
|
272
|
+
|
|
273
|
+
_ENV_PORT = "PORT"
|
|
274
|
+
_ENV_STUDIO_ORIGIN = "STUDIO_API_ORIGIN"
|
|
275
|
+
_ENV_LEASE_AUDIENCE = "MR_SESSION_WORKER_LEASE_AUDIENCE"
|
|
276
|
+
_ENV_WORKER_BOOTSTRAP_AUDIENCE = "STUDIO_WORKER_BOOTSTRAP_AUDIENCE"
|
|
277
|
+
_ENV_SHUTDOWN_GRACE = "MR_SESSION_WORKER_SHUTDOWN_GRACE_SECONDS"
|
|
278
|
+
_ENV_WARM_POOL = "MR_SESSION_WORKER_WARM_POOL_SIZE"
|
|
279
|
+
_ENV_PROBE_TIMEOUT = "MR_SESSION_WORKER_PROBE_TIMEOUT_SECONDS"
|
|
280
|
+
_ENV_ACQUISITION_TIMEOUT = "MR_SESSION_WORKER_ACQUISITION_TIMEOUT_SECONDS"
|
|
281
|
+
_ENV_EXTERNAL_NETWORK_POLICY = "MOSTLYRIGHT_EXTERNAL_NETWORK_POLICY_ATTESTATION"
|
|
282
|
+
_ENV_CRAWLER_RESULTS_BUCKET = "MOSTLYRIGHT_CRAWLER_RESULTS_BUCKET"
|
|
283
|
+
|
|
284
|
+
|
|
285
|
+
class SessionWorkerError(RuntimeError):
|
|
286
|
+
"""A coded session-worker refusal. The code is what a probe or a session fails with."""
|
|
287
|
+
|
|
288
|
+
def __init__(self, code: str, detail: str, *, retry_after_seconds: int | None = None) -> None:
|
|
289
|
+
super().__init__(f"{code}: {detail}")
|
|
290
|
+
self.code = code
|
|
291
|
+
self.detail = detail
|
|
292
|
+
self.retry_after_seconds = retry_after_seconds
|
|
293
|
+
|
|
294
|
+
|
|
295
|
+
@dataclass(frozen=True)
|
|
296
|
+
class SessionDispatch:
|
|
297
|
+
"""The complete, non-secret argument contract of one Cloud Run Job task.
|
|
298
|
+
|
|
299
|
+
Nothing from the resulting lease belongs here. In particular, keeping this object safe to
|
|
300
|
+
log is a guard against a future Terraform change accidentally moving a capability or nonce
|
|
301
|
+
into command arguments, which Cloud Run exposes as execution configuration.
|
|
302
|
+
"""
|
|
303
|
+
|
|
304
|
+
dispatch_id: str
|
|
305
|
+
exchange_url: str
|
|
306
|
+
audience: str
|
|
307
|
+
|
|
308
|
+
|
|
309
|
+
# ==========================================================================================
|
|
310
|
+
# Configuration
|
|
311
|
+
# ==========================================================================================
|
|
312
|
+
|
|
313
|
+
|
|
314
|
+
@dataclass(frozen=True)
|
|
315
|
+
class SessionWorkerConfig:
|
|
316
|
+
"""Everything the service reads from its environment, validated once at start."""
|
|
317
|
+
|
|
318
|
+
studio_origin: str
|
|
319
|
+
lease_audience: str
|
|
320
|
+
port: int = 8080
|
|
321
|
+
shutdown_grace_seconds: float = DEFAULT_SHUTDOWN_GRACE_SECONDS
|
|
322
|
+
warm_pool_size: int = DEFAULT_WARM_POOL_SIZE
|
|
323
|
+
probe_timeout_seconds: float = DEFAULT_PROBE_TIMEOUT_SECONDS
|
|
324
|
+
acquisition_timeout_seconds: float = DEFAULT_ACQUISITION_TIMEOUT_SECONDS
|
|
325
|
+
external_network_policy_attestation: str | None = None
|
|
326
|
+
crawler_results_bucket: str | None = None
|
|
327
|
+
|
|
328
|
+
|
|
329
|
+
def _positive_number(raw: str | None, default: float, label: str, *, maximum: float) -> float:
|
|
330
|
+
if raw is None or raw == "":
|
|
331
|
+
return default
|
|
332
|
+
try:
|
|
333
|
+
value = float(raw)
|
|
334
|
+
except ValueError as error:
|
|
335
|
+
raise SessionWorkerError("SESSION_WORKER_CONFIG_INVALID", f"{label} is not a number") from (
|
|
336
|
+
error
|
|
337
|
+
)
|
|
338
|
+
if not 0 < value <= maximum:
|
|
339
|
+
raise SessionWorkerError(
|
|
340
|
+
"SESSION_WORKER_CONFIG_INVALID", f"{label} must be in (0, {maximum}]"
|
|
341
|
+
)
|
|
342
|
+
return value
|
|
343
|
+
|
|
344
|
+
|
|
345
|
+
def config_from_environment(environment: Mapping[str, str] | None = None) -> SessionWorkerConfig:
|
|
346
|
+
"""Build the config, refusing rather than defaulting anything that is a coordinate."""
|
|
347
|
+
|
|
348
|
+
env = os.environ if environment is None else environment
|
|
349
|
+
origin = env.get(_ENV_STUDIO_ORIGIN, "")
|
|
350
|
+
if not origin:
|
|
351
|
+
raise SessionWorkerError(
|
|
352
|
+
"SESSION_WORKER_CONFIG_INVALID", f"{_ENV_STUDIO_ORIGIN} is required"
|
|
353
|
+
)
|
|
354
|
+
origin = origin.rstrip("/")
|
|
355
|
+
_https_url(origin, _ENV_STUDIO_ORIGIN)
|
|
356
|
+
# Studio verifies the session worker against the same closed workload-identity audience as the
|
|
357
|
+
# producer bootstrap exchange. The shared Cloud Run template already supplies that coordinate;
|
|
358
|
+
# consuming it here keeps the identity token's ``aud`` exact while preserving the explicit
|
|
359
|
+
# Harness override for a deliberately split deployment. An older deployment without the shared
|
|
360
|
+
# coordinate remains compatible with its historical origin audience.
|
|
361
|
+
audience = (
|
|
362
|
+
env.get(_ENV_LEASE_AUDIENCE, "") or env.get(_ENV_WORKER_BOOTSTRAP_AUDIENCE, "") or origin
|
|
363
|
+
)
|
|
364
|
+
_https_url(
|
|
365
|
+
audience,
|
|
366
|
+
(
|
|
367
|
+
_ENV_LEASE_AUDIENCE
|
|
368
|
+
if env.get(_ENV_LEASE_AUDIENCE, "")
|
|
369
|
+
else _ENV_WORKER_BOOTSTRAP_AUDIENCE
|
|
370
|
+
if env.get(_ENV_WORKER_BOOTSTRAP_AUDIENCE, "")
|
|
371
|
+
else _ENV_STUDIO_ORIGIN
|
|
372
|
+
),
|
|
373
|
+
)
|
|
374
|
+
raw_port = env.get(_ENV_PORT, "8080")
|
|
375
|
+
try:
|
|
376
|
+
port = int(raw_port)
|
|
377
|
+
except ValueError as error:
|
|
378
|
+
raise SessionWorkerError("SESSION_WORKER_CONFIG_INVALID", "PORT is not an integer") from (
|
|
379
|
+
error
|
|
380
|
+
)
|
|
381
|
+
if not 1 <= port <= 65535:
|
|
382
|
+
raise SessionWorkerError("SESSION_WORKER_CONFIG_INVALID", "PORT is out of range")
|
|
383
|
+
raw_pool = env.get(_ENV_WARM_POOL, "")
|
|
384
|
+
try:
|
|
385
|
+
warm_pool = DEFAULT_WARM_POOL_SIZE if raw_pool == "" else int(raw_pool)
|
|
386
|
+
except ValueError as error:
|
|
387
|
+
raise SessionWorkerError(
|
|
388
|
+
"SESSION_WORKER_CONFIG_INVALID", f"{_ENV_WARM_POOL} is not an integer"
|
|
389
|
+
) from error
|
|
390
|
+
if not 0 <= warm_pool <= 8:
|
|
391
|
+
raise SessionWorkerError("SESSION_WORKER_CONFIG_INVALID", f"{_ENV_WARM_POOL} must be 0..8")
|
|
392
|
+
return SessionWorkerConfig(
|
|
393
|
+
studio_origin=origin,
|
|
394
|
+
lease_audience=audience,
|
|
395
|
+
port=port,
|
|
396
|
+
shutdown_grace_seconds=_positive_number(
|
|
397
|
+
env.get(_ENV_SHUTDOWN_GRACE),
|
|
398
|
+
DEFAULT_SHUTDOWN_GRACE_SECONDS,
|
|
399
|
+
_ENV_SHUTDOWN_GRACE,
|
|
400
|
+
maximum=600.0,
|
|
401
|
+
),
|
|
402
|
+
warm_pool_size=warm_pool,
|
|
403
|
+
probe_timeout_seconds=_positive_number(
|
|
404
|
+
env.get(_ENV_PROBE_TIMEOUT),
|
|
405
|
+
DEFAULT_PROBE_TIMEOUT_SECONDS,
|
|
406
|
+
_ENV_PROBE_TIMEOUT,
|
|
407
|
+
maximum=300.0,
|
|
408
|
+
),
|
|
409
|
+
acquisition_timeout_seconds=_positive_number(
|
|
410
|
+
env.get(_ENV_ACQUISITION_TIMEOUT),
|
|
411
|
+
DEFAULT_ACQUISITION_TIMEOUT_SECONDS,
|
|
412
|
+
_ENV_ACQUISITION_TIMEOUT,
|
|
413
|
+
maximum=600.0,
|
|
414
|
+
),
|
|
415
|
+
external_network_policy_attestation=env.get(_ENV_EXTERNAL_NETWORK_POLICY) or None,
|
|
416
|
+
# The same name the hosted producer reads, because it is the same quarantine bucket and a
|
|
417
|
+
# second spelling would be a second thing to get wrong on one deployment. Absent means a
|
|
418
|
+
# Studio-local result authority, which `validate_result_download` then requires.
|
|
419
|
+
crawler_results_bucket=env.get(_ENV_CRAWLER_RESULTS_BUCKET) or None,
|
|
420
|
+
)
|
|
421
|
+
|
|
422
|
+
|
|
423
|
+
# ==========================================================================================
|
|
424
|
+
# Cloud Run Job dispatch bootstrap
|
|
425
|
+
# ==========================================================================================
|
|
426
|
+
|
|
427
|
+
|
|
428
|
+
def _dispatch_uuid(value: Any) -> str:
|
|
429
|
+
if not isinstance(value, str):
|
|
430
|
+
raise SessionWorkerError("SESSION_DISPATCH_INVALID", "dispatch id must be a string")
|
|
431
|
+
try:
|
|
432
|
+
parsed = UUID(value)
|
|
433
|
+
except ValueError as error:
|
|
434
|
+
raise SessionWorkerError("SESSION_DISPATCH_INVALID", "dispatch id is not a UUID") from error
|
|
435
|
+
if parsed.version != 4 or str(parsed) != value:
|
|
436
|
+
raise SessionWorkerError(
|
|
437
|
+
"SESSION_DISPATCH_INVALID", "dispatch id must be a canonical UUID4"
|
|
438
|
+
)
|
|
439
|
+
return value
|
|
440
|
+
|
|
441
|
+
|
|
442
|
+
def _dispatch_https_url(value: Any, label: str) -> str:
|
|
443
|
+
try:
|
|
444
|
+
return _https_url(value, label)
|
|
445
|
+
except HostedWorkerError as error:
|
|
446
|
+
raise SessionWorkerError("SESSION_DISPATCH_INVALID", error.detail) from error
|
|
447
|
+
|
|
448
|
+
|
|
449
|
+
def parse_session_dispatch(argv: Sequence[str]) -> SessionDispatch:
|
|
450
|
+
"""Parse exactly the three non-secret Job arguments, in their deployment order.
|
|
451
|
+
|
|
452
|
+
A normal argument parser intentionally permits flags in any order and can silently admit
|
|
453
|
+
future options. Job overrides are a security boundary, so this parser instead spells the
|
|
454
|
+
only six argv cells Studio may create. It also keeps the invocation object incapable of
|
|
455
|
+
carrying a lease nonce or capability.
|
|
456
|
+
"""
|
|
457
|
+
|
|
458
|
+
values = tuple(argv)
|
|
459
|
+
flags = (
|
|
460
|
+
"--session-dispatch-id",
|
|
461
|
+
"--session-dispatch-exchange-url",
|
|
462
|
+
"--session-dispatch-audience",
|
|
463
|
+
)
|
|
464
|
+
if len(values) != 6 or values[::2] != flags or any(not item for item in values[1::2]):
|
|
465
|
+
raise SessionWorkerError(
|
|
466
|
+
"SESSION_DISPATCH_ARGUMENTS_INVALID",
|
|
467
|
+
"expected --session-dispatch-id, --session-dispatch-exchange-url and "
|
|
468
|
+
"--session-dispatch-audience in that order",
|
|
469
|
+
)
|
|
470
|
+
dispatch_id, exchange_url, audience = values[1::2]
|
|
471
|
+
return SessionDispatch(
|
|
472
|
+
dispatch_id=_dispatch_uuid(dispatch_id),
|
|
473
|
+
exchange_url=_dispatch_https_url(exchange_url, "session dispatch exchange URL"),
|
|
474
|
+
audience=_dispatch_https_url(audience, "session dispatch audience"),
|
|
475
|
+
)
|
|
476
|
+
|
|
477
|
+
|
|
478
|
+
def _validate_session_dispatch(config: SessionWorkerConfig, dispatch: SessionDispatch) -> None:
|
|
479
|
+
"""Bind a Job's opaque id to this deployment's one exchange endpoint.
|
|
480
|
+
|
|
481
|
+
The OIDC audience is intentionally supplied by Studio's Job override, but the URL itself must
|
|
482
|
+
remain derived from the locally pinned Studio origin. This closes request-target substitution
|
|
483
|
+
even if an operator can create an execution override.
|
|
484
|
+
"""
|
|
485
|
+
|
|
486
|
+
expected = config.studio_origin + SESSION_DISPATCH_EXCHANGE_PATH_TEMPLATE.format(
|
|
487
|
+
dispatch_id=dispatch.dispatch_id
|
|
488
|
+
)
|
|
489
|
+
if dispatch.exchange_url != expected:
|
|
490
|
+
raise SessionWorkerError(
|
|
491
|
+
"SESSION_DISPATCH_URL_UNRECOGNIZED",
|
|
492
|
+
"session dispatch exchange URL is not this deployment's endpoint for this dispatch",
|
|
493
|
+
)
|
|
494
|
+
if dispatch.audience != config.studio_origin:
|
|
495
|
+
raise SessionWorkerError(
|
|
496
|
+
"SESSION_DISPATCH_AUDIENCE_UNRECOGNIZED",
|
|
497
|
+
"session dispatch audience is not this deployment's Studio audience",
|
|
498
|
+
)
|
|
499
|
+
|
|
500
|
+
|
|
501
|
+
def _dispatch_exchange_error(status: int, raw: bytearray) -> SessionWorkerError:
|
|
502
|
+
"""Keep Studio's typed refusal without echoing a response that could contain a secret."""
|
|
503
|
+
|
|
504
|
+
try:
|
|
505
|
+
value = _strict_json(raw, "session dispatch exchange refusal")
|
|
506
|
+
except HostedWorkerError:
|
|
507
|
+
value = None
|
|
508
|
+
if isinstance(value, dict):
|
|
509
|
+
code = value.get("code")
|
|
510
|
+
if isinstance(code, str) and re.fullmatch(r"[A-Z][A-Z0-9_]{2,127}", code):
|
|
511
|
+
return SessionWorkerError(code, f"session dispatch exchange returned HTTP {status}")
|
|
512
|
+
return SessionWorkerError(
|
|
513
|
+
"SESSION_DISPATCH_EXCHANGE_FAILED", f"session dispatch exchange returned HTTP {status}"
|
|
514
|
+
)
|
|
515
|
+
|
|
516
|
+
|
|
517
|
+
def exchange_session_dispatch(
|
|
518
|
+
*,
|
|
519
|
+
config: SessionWorkerConfig,
|
|
520
|
+
dispatch: SessionDispatch,
|
|
521
|
+
transport: BootstrapTransport | None = None,
|
|
522
|
+
) -> tuple[SessionLease, CredentialLease]:
|
|
523
|
+
"""Redeem one dispatch through Studio and return the sole fenced session lease.
|
|
524
|
+
|
|
525
|
+
The metadata identity proves the Job's workload principal. The body is deliberately exact and
|
|
526
|
+
contains no secret: Studio's atomically claimed dispatch record is what binds that principal to
|
|
527
|
+
the resulting attempt. An already claimed or terminal dispatch is a typed non-200 response;
|
|
528
|
+
callers must treat it as terminal rather than retrying into a second supervisor.
|
|
529
|
+
"""
|
|
530
|
+
|
|
531
|
+
_validate_session_dispatch(config, dispatch)
|
|
532
|
+
active_transport = transport or UrlLibBootstrapTransport()
|
|
533
|
+
identity = _metadata_identity(active_transport, dispatch.audience)
|
|
534
|
+
payload = canonical_json_bytes(
|
|
535
|
+
{
|
|
536
|
+
"schema_version": SESSION_DISPATCH_SCHEMA_VERSION,
|
|
537
|
+
"dispatch_id": dispatch.dispatch_id,
|
|
538
|
+
}
|
|
539
|
+
)
|
|
540
|
+
try:
|
|
541
|
+
token = identity.decode("ascii")
|
|
542
|
+
status, response = active_transport.request(
|
|
543
|
+
"POST",
|
|
544
|
+
dispatch.exchange_url,
|
|
545
|
+
{
|
|
546
|
+
"Authorization": f"Bearer {token}",
|
|
547
|
+
"Content-Type": "application/json",
|
|
548
|
+
},
|
|
549
|
+
payload,
|
|
550
|
+
MAX_SESSION_DISPATCH_RESPONSE_BYTES,
|
|
551
|
+
)
|
|
552
|
+
finally:
|
|
553
|
+
identity[:] = b"\x00" * len(identity)
|
|
554
|
+
raw = bytearray(response)
|
|
555
|
+
try:
|
|
556
|
+
if status != HTTPStatus.OK:
|
|
557
|
+
raise _dispatch_exchange_error(status, raw)
|
|
558
|
+
try:
|
|
559
|
+
lease_wire = _strict_json(raw, "session dispatch exchange response")
|
|
560
|
+
except HostedWorkerError as error:
|
|
561
|
+
raise SessionWorkerError(
|
|
562
|
+
"SESSION_DISPATCH_RESPONSE_INVALID", "session dispatch exchange response is invalid"
|
|
563
|
+
) from error
|
|
564
|
+
try:
|
|
565
|
+
return parse_session_lease(lease_wire)
|
|
566
|
+
finally:
|
|
567
|
+
if isinstance(lease_wire, dict):
|
|
568
|
+
lease_wire.clear()
|
|
569
|
+
finally:
|
|
570
|
+
raw[:] = b"\x00" * len(raw)
|
|
571
|
+
|
|
572
|
+
|
|
573
|
+
# ==========================================================================================
|
|
574
|
+
# The start invocation
|
|
575
|
+
# ==========================================================================================
|
|
576
|
+
|
|
577
|
+
|
|
578
|
+
@dataclass(frozen=True)
|
|
579
|
+
class StartInvocation:
|
|
580
|
+
"""The four non-secret coordinates of one start signal. The nonce is never in here."""
|
|
581
|
+
|
|
582
|
+
session_id: str
|
|
583
|
+
workspace_id: str
|
|
584
|
+
lease_url: str
|
|
585
|
+
|
|
586
|
+
|
|
587
|
+
def _uuid_text(value: Any, label: str) -> str:
|
|
588
|
+
if not isinstance(value, str):
|
|
589
|
+
raise SessionWorkerError("SESSION_START_INVALID", f"{label} must be a string")
|
|
590
|
+
try:
|
|
591
|
+
parsed = UUID(value)
|
|
592
|
+
except ValueError as error:
|
|
593
|
+
raise SessionWorkerError("SESSION_START_INVALID", f"{label} is not a UUID") from error
|
|
594
|
+
if str(parsed) != value:
|
|
595
|
+
raise SessionWorkerError("SESSION_START_INVALID", f"{label} is not canonical")
|
|
596
|
+
return value
|
|
597
|
+
|
|
598
|
+
|
|
599
|
+
def parse_start_invocation(raw: bytes, *, studio_origin: str) -> tuple[StartInvocation, bytearray]:
|
|
600
|
+
"""Admit one start invocation and separate the nonce from everything else.
|
|
601
|
+
|
|
602
|
+
The nonce comes back as a ``bytearray`` rather than a ``str`` so the caller can zero it, and it
|
|
603
|
+
is never a field on :class:`StartInvocation` so that logging or repr'ing the invocation cannot
|
|
604
|
+
print it.
|
|
605
|
+
|
|
606
|
+
``lease_url`` is **checked, never followed**. The URL arrives inside a request body, and a
|
|
607
|
+
request body is data: a worker that fetched whatever URL it was handed -- carrying a Google
|
|
608
|
+
identity token minted for the audience it was configured with -- would be an SSRF primitive
|
|
609
|
+
with a credential attached, reachable by anything that could reach this port. The value is
|
|
610
|
+
therefore compared against the URL derived from ``STUDIO_API_ORIGIN`` and the session id, and a
|
|
611
|
+
mismatch is a refusal. The call itself is then made through the generated client bound to
|
|
612
|
+
``STUDIO_API_ORIGIN``, so even a check that somehow passed could not redirect the request.
|
|
613
|
+
"""
|
|
614
|
+
|
|
615
|
+
if len(raw) > MAX_START_REQUEST_BYTES:
|
|
616
|
+
raise SessionWorkerError("SESSION_START_TOO_LARGE", "start invocation exceeds its budget")
|
|
617
|
+
try:
|
|
618
|
+
value = json.loads(raw.decode("utf-8"))
|
|
619
|
+
except (UnicodeDecodeError, ValueError) as error:
|
|
620
|
+
raise SessionWorkerError("SESSION_START_INVALID", "start invocation is not JSON") from error
|
|
621
|
+
if not isinstance(value, dict):
|
|
622
|
+
raise SessionWorkerError("SESSION_START_INVALID", "start invocation must be an object")
|
|
623
|
+
expected = {"schema_version", "session_id", "workspace_id", "lease_nonce", "lease_url"}
|
|
624
|
+
if set(value) != expected:
|
|
625
|
+
raise SessionWorkerError("SESSION_START_INVALID", "start invocation fields are not exact")
|
|
626
|
+
if value["schema_version"] != START_SCHEMA_VERSION:
|
|
627
|
+
raise SessionWorkerError("SESSION_START_INVALID", "start invocation schema is unsupported")
|
|
628
|
+
session_id = _uuid_text(value["session_id"], "session_id")
|
|
629
|
+
workspace_id = _uuid_text(value["workspace_id"], "workspace_id")
|
|
630
|
+
nonce = value["lease_nonce"]
|
|
631
|
+
if (
|
|
632
|
+
not isinstance(nonce, str)
|
|
633
|
+
or not MIN_NONCE_CHARS <= len(nonce) <= MAX_NONCE_CHARS
|
|
634
|
+
or not all(character.isalnum() or character in "_-" for character in nonce)
|
|
635
|
+
):
|
|
636
|
+
raise SessionWorkerError("SESSION_START_INVALID", "lease nonce shape is invalid")
|
|
637
|
+
lease_url = value["lease_url"]
|
|
638
|
+
if not isinstance(lease_url, str):
|
|
639
|
+
raise SessionWorkerError("SESSION_START_INVALID", "lease_url must be a string")
|
|
640
|
+
derived = studio_origin.rstrip("/") + LEASE_PATH_TEMPLATE.format(session_id=session_id)
|
|
641
|
+
if lease_url != derived:
|
|
642
|
+
raise SessionWorkerError(
|
|
643
|
+
"SESSION_LEASE_URL_UNRECOGNIZED",
|
|
644
|
+
"lease_url is not this deployment's lease endpoint for this session",
|
|
645
|
+
)
|
|
646
|
+
return (
|
|
647
|
+
StartInvocation(session_id=session_id, workspace_id=workspace_id, lease_url=derived),
|
|
648
|
+
bytearray(nonce.encode("ascii")),
|
|
649
|
+
)
|
|
650
|
+
|
|
651
|
+
|
|
652
|
+
# ==========================================================================================
|
|
653
|
+
# The lease
|
|
654
|
+
# ==========================================================================================
|
|
655
|
+
|
|
656
|
+
|
|
657
|
+
@dataclass(frozen=True)
|
|
658
|
+
class SessionContext:
|
|
659
|
+
"""The workspace coordinate the session probes against. Also the probe authority allowlist."""
|
|
660
|
+
|
|
661
|
+
dataset_id: str
|
|
662
|
+
question_id: str
|
|
663
|
+
table_id: str | None
|
|
664
|
+
source_ids: frozenset[str]
|
|
665
|
+
|
|
666
|
+
|
|
667
|
+
@dataclass(frozen=True)
|
|
668
|
+
class SessionLease:
|
|
669
|
+
"""Everything the lease said except the capability, which lives in a ``CredentialLease``."""
|
|
670
|
+
|
|
671
|
+
session_id: str
|
|
672
|
+
workspace_id: UUID
|
|
673
|
+
run_id: UUID
|
|
674
|
+
attempt_id: UUID
|
|
675
|
+
generation: int
|
|
676
|
+
fence: int
|
|
677
|
+
capability_expires_at: datetime
|
|
678
|
+
callback_base_url: str
|
|
679
|
+
allowed_callback_paths: tuple[str, ...]
|
|
680
|
+
heartbeat_interval_seconds: float
|
|
681
|
+
claim_poll_interval_seconds: float
|
|
682
|
+
idle_timeout_seconds: float
|
|
683
|
+
context: SessionContext
|
|
684
|
+
|
|
685
|
+
|
|
686
|
+
def _positive_int(value: Any, label: str) -> int:
|
|
687
|
+
if type(value) is not int or value <= 0 or value > 2**53 - 1:
|
|
688
|
+
raise SessionWorkerError("SESSION_LEASE_INVALID", f"{label} is not a positive integer")
|
|
689
|
+
return value
|
|
690
|
+
|
|
691
|
+
|
|
692
|
+
def _parse_utc_timestamp(value: Any, label: str) -> datetime:
|
|
693
|
+
"""Parse Studio's canonical UTC instant without accepting a local-time interpretation."""
|
|
694
|
+
|
|
695
|
+
if not isinstance(value, str) or _UTC_TIMESTAMP.fullmatch(value) is None:
|
|
696
|
+
raise SessionWorkerError(
|
|
697
|
+
"SESSION_LEASE_INVALID", f"{label} must be a canonical UTC timestamp"
|
|
698
|
+
)
|
|
699
|
+
try:
|
|
700
|
+
pattern = "%Y-%m-%dT%H:%M:%S.%fZ" if "." in value else "%Y-%m-%dT%H:%M:%SZ"
|
|
701
|
+
return datetime.strptime(value, pattern).replace(tzinfo=UTC)
|
|
702
|
+
except ValueError as error:
|
|
703
|
+
raise SessionWorkerError(
|
|
704
|
+
"SESSION_LEASE_INVALID", f"{label} is not a real UTC timestamp"
|
|
705
|
+
) from error
|
|
706
|
+
|
|
707
|
+
|
|
708
|
+
def _wall_clock_now() -> datetime:
|
|
709
|
+
return datetime.now(UTC)
|
|
710
|
+
|
|
711
|
+
|
|
712
|
+
@dataclass(frozen=True)
|
|
713
|
+
class SessionDeadline:
|
|
714
|
+
"""One capability boundary measured against both wall and monotonic clocks.
|
|
715
|
+
|
|
716
|
+
The wire instant is the authority, but a wall clock can jump backwards after the lease is
|
|
717
|
+
accepted. Capturing its remaining duration on a monotonic clock prevents that jump from
|
|
718
|
+
extending a session. Conversely, a forward wall-clock jump stops early, which is the safe
|
|
719
|
+
direction for a fenced credential.
|
|
720
|
+
"""
|
|
721
|
+
|
|
722
|
+
expires_at: datetime
|
|
723
|
+
monotonic_expires_at: float
|
|
724
|
+
_wall_clock: Callable[[], datetime] = field(repr=False, compare=False)
|
|
725
|
+
_monotonic_clock: Callable[[], float] = field(repr=False, compare=False)
|
|
726
|
+
|
|
727
|
+
@classmethod
|
|
728
|
+
def from_lease(
|
|
729
|
+
cls,
|
|
730
|
+
lease: SessionLease,
|
|
731
|
+
*,
|
|
732
|
+
wall_clock: Callable[[], datetime] = _wall_clock_now,
|
|
733
|
+
monotonic_clock: Callable[[], float] = time.monotonic,
|
|
734
|
+
) -> SessionDeadline:
|
|
735
|
+
return cls.from_expiry(
|
|
736
|
+
lease.capability_expires_at,
|
|
737
|
+
wall_clock=wall_clock,
|
|
738
|
+
monotonic_clock=monotonic_clock,
|
|
739
|
+
)
|
|
740
|
+
|
|
741
|
+
@classmethod
|
|
742
|
+
def from_expiry(
|
|
743
|
+
cls,
|
|
744
|
+
expires_at: datetime,
|
|
745
|
+
*,
|
|
746
|
+
wall_clock: Callable[[], datetime] = _wall_clock_now,
|
|
747
|
+
monotonic_clock: Callable[[], float] = time.monotonic,
|
|
748
|
+
) -> SessionDeadline:
|
|
749
|
+
"""Create a paired deadline from an already authenticated expiry instant."""
|
|
750
|
+
|
|
751
|
+
if not isinstance(expires_at, datetime) or expires_at.tzinfo is None:
|
|
752
|
+
raise SessionWorkerError(
|
|
753
|
+
"SESSION_LEASE_INVALID", "capability_expires_at is not an absolute instant"
|
|
754
|
+
)
|
|
755
|
+
expires_at = expires_at.astimezone(UTC)
|
|
756
|
+
now = cls._wall_now(wall_clock)
|
|
757
|
+
remaining = (expires_at - now).total_seconds()
|
|
758
|
+
if remaining <= 0:
|
|
759
|
+
raise SessionWorkerError(
|
|
760
|
+
"SESSION_CAPABILITY_EXPIRED",
|
|
761
|
+
"the session capability expired before the worker could use it",
|
|
762
|
+
)
|
|
763
|
+
if remaining > MAX_SESSION_CAPABILITY_REMAINING_SECONDS:
|
|
764
|
+
raise SessionWorkerError(
|
|
765
|
+
"SESSION_LEASE_EXPIRY_UNREASONABLE",
|
|
766
|
+
"the session capability exceeds the maximum producer lease",
|
|
767
|
+
)
|
|
768
|
+
monotonic_now = cls._monotonic_now(monotonic_clock)
|
|
769
|
+
return cls(
|
|
770
|
+
expires_at=expires_at,
|
|
771
|
+
monotonic_expires_at=monotonic_now + remaining,
|
|
772
|
+
_wall_clock=wall_clock,
|
|
773
|
+
_monotonic_clock=monotonic_clock,
|
|
774
|
+
)
|
|
775
|
+
|
|
776
|
+
@staticmethod
|
|
777
|
+
def _wall_now(clock: Callable[[], datetime]) -> datetime:
|
|
778
|
+
now = clock()
|
|
779
|
+
if not isinstance(now, datetime) or now.tzinfo is None:
|
|
780
|
+
raise SessionWorkerError(
|
|
781
|
+
"SESSION_CLOCK_INVALID", "the session worker wall clock is not an absolute instant"
|
|
782
|
+
)
|
|
783
|
+
return now.astimezone(UTC)
|
|
784
|
+
|
|
785
|
+
@staticmethod
|
|
786
|
+
def _monotonic_now(clock: Callable[[], float]) -> float:
|
|
787
|
+
value = float(clock())
|
|
788
|
+
if not math.isfinite(value):
|
|
789
|
+
raise SessionWorkerError(
|
|
790
|
+
"SESSION_CLOCK_INVALID", "the session worker monotonic clock is not finite"
|
|
791
|
+
)
|
|
792
|
+
return value
|
|
793
|
+
|
|
794
|
+
def remaining_seconds(self) -> float:
|
|
795
|
+
"""Return the earliest safe deadline; a broken clock fails closed at zero."""
|
|
796
|
+
|
|
797
|
+
try:
|
|
798
|
+
wall_remaining = (self.expires_at - self._wall_now(self._wall_clock)).total_seconds()
|
|
799
|
+
monotonic_remaining = self.monotonic_expires_at - self._monotonic_now(
|
|
800
|
+
self._monotonic_clock
|
|
801
|
+
)
|
|
802
|
+
except (SessionWorkerError, TypeError, ValueError, OverflowError):
|
|
803
|
+
return 0.0
|
|
804
|
+
return min(wall_remaining, monotonic_remaining)
|
|
805
|
+
|
|
806
|
+
def expired(self) -> bool:
|
|
807
|
+
return self.remaining_seconds() <= 0
|
|
808
|
+
|
|
809
|
+
def bounded_wait(self, requested_seconds: float) -> float:
|
|
810
|
+
"""Limit one interruptible wait to the remaining credential lifetime."""
|
|
811
|
+
|
|
812
|
+
return max(0.0, min(float(requested_seconds), self.remaining_seconds()))
|
|
813
|
+
|
|
814
|
+
|
|
815
|
+
def _linux_direct_child_pids() -> frozenset[int]:
|
|
816
|
+
"""Return direct children across threads, or fail closed if Linux cannot enumerate them."""
|
|
817
|
+
|
|
818
|
+
if sys.platform != "linux":
|
|
819
|
+
return frozenset()
|
|
820
|
+
children: set[int] = set()
|
|
821
|
+
try:
|
|
822
|
+
for task in Path("/proc/self/task").iterdir():
|
|
823
|
+
raw = (task / "children").read_text(encoding="ascii").strip()
|
|
824
|
+
if raw:
|
|
825
|
+
children.update(int(value) for value in raw.split())
|
|
826
|
+
except (OSError, ValueError) as error:
|
|
827
|
+
raise SessionWorkerError(
|
|
828
|
+
"SESSION_PROCESS_ISOLATION_UNAVAILABLE",
|
|
829
|
+
"could not enumerate isolated-session descendants",
|
|
830
|
+
) from error
|
|
831
|
+
return frozenset(children)
|
|
832
|
+
|
|
833
|
+
|
|
834
|
+
def parse_session_lease(value: Any) -> tuple[SessionLease, CredentialLease]:
|
|
835
|
+
"""Admit one ``ResearchSessionLease`` and put the capability straight into a lease object.
|
|
836
|
+
|
|
837
|
+
The capability never becomes a module-level or dataclass-held ``str``. It is validated by
|
|
838
|
+
``hosted_worker.parse_capability`` -- the same shape check the producer path applies -- and then
|
|
839
|
+
handed to ``CredentialLease``, which refuses to be pickled, refuses ``to_dict``, and renders as
|
|
840
|
+
``<CredentialLease redacted state=open>`` wherever anything tries to print it.
|
|
841
|
+
"""
|
|
842
|
+
|
|
843
|
+
if not isinstance(value, dict):
|
|
844
|
+
raise SessionWorkerError("SESSION_LEASE_INVALID", "lease must be an object")
|
|
845
|
+
required = {
|
|
846
|
+
"schema_version",
|
|
847
|
+
"session_id",
|
|
848
|
+
"workspace_id",
|
|
849
|
+
"run_id",
|
|
850
|
+
"attempt_id",
|
|
851
|
+
"generation",
|
|
852
|
+
"fence",
|
|
853
|
+
"capability",
|
|
854
|
+
"capability_expires_at",
|
|
855
|
+
"callback_base_url",
|
|
856
|
+
"allowed_callback_paths",
|
|
857
|
+
"heartbeat_interval_seconds",
|
|
858
|
+
"claim_poll_interval_seconds",
|
|
859
|
+
"idle_timeout_seconds",
|
|
860
|
+
"context",
|
|
861
|
+
}
|
|
862
|
+
if set(value) != required:
|
|
863
|
+
missing = sorted(required - set(value))
|
|
864
|
+
unexpected = sorted(set(value) - required)
|
|
865
|
+
detail = []
|
|
866
|
+
if missing:
|
|
867
|
+
detail.append(f"omits {', '.join(missing)}")
|
|
868
|
+
if unexpected:
|
|
869
|
+
detail.append(f"has unexpected {', '.join(unexpected)}")
|
|
870
|
+
raise SessionWorkerError("SESSION_LEASE_INVALID", f"lease {'; '.join(detail)}")
|
|
871
|
+
expires_at = _parse_utc_timestamp(value["capability_expires_at"], "capability_expires_at")
|
|
872
|
+
raw_capability = value["capability"]
|
|
873
|
+
if not isinstance(raw_capability, str) or len(raw_capability) > MAX_CAPABILITY_BYTES:
|
|
874
|
+
raise SessionWorkerError("SESSION_LEASE_INVALID", "lease capability is not admissible")
|
|
875
|
+
# `parse_capability` refuses anything that is not the exact `mr_cap_...` shape, so a lease that
|
|
876
|
+
# somehow carried a URL, a path, or a JWT never reaches an Authorization header.
|
|
877
|
+
token = parse_capability(raw_capability.encode("ascii", errors="replace"))
|
|
878
|
+
credential = CredentialLease(token.encode("ascii"), {"kind": "session-attempt-capability"})
|
|
879
|
+
del token, raw_capability
|
|
880
|
+
|
|
881
|
+
paths = value["allowed_callback_paths"]
|
|
882
|
+
if not isinstance(paths, list) or not paths:
|
|
883
|
+
credential.close()
|
|
884
|
+
raise SessionWorkerError("SESSION_LEASE_INVALID", "lease declares no callback paths")
|
|
885
|
+
for entry in paths:
|
|
886
|
+
if not isinstance(entry, str) or not entry.startswith("/internal/v3/producer/"):
|
|
887
|
+
credential.close()
|
|
888
|
+
raise SessionWorkerError(
|
|
889
|
+
"SESSION_LEASE_INVALID", "lease callback path is outside the Builder surface"
|
|
890
|
+
)
|
|
891
|
+
context = value["context"]
|
|
892
|
+
if not isinstance(context, dict):
|
|
893
|
+
credential.close()
|
|
894
|
+
raise SessionWorkerError("SESSION_LEASE_INVALID", "lease context must be an object")
|
|
895
|
+
source_ids = context.get("source_ids", [])
|
|
896
|
+
if not isinstance(source_ids, list) or len(source_ids) > 32:
|
|
897
|
+
credential.close()
|
|
898
|
+
raise SessionWorkerError("SESSION_LEASE_INVALID", "lease context source_ids is invalid")
|
|
899
|
+
try:
|
|
900
|
+
base_url = value["callback_base_url"]
|
|
901
|
+
_https_url(str(base_url).rstrip("/"), "callback_base_url")
|
|
902
|
+
lease = SessionLease(
|
|
903
|
+
session_id=_uuid_text(value["session_id"], "session_id"),
|
|
904
|
+
workspace_id=UUID(_uuid_text(value["workspace_id"], "workspace_id")),
|
|
905
|
+
run_id=UUID(_uuid_text(value["run_id"], "run_id")),
|
|
906
|
+
attempt_id=UUID(_uuid_text(value["attempt_id"], "attempt_id")),
|
|
907
|
+
generation=_positive_int(value["generation"], "generation"),
|
|
908
|
+
fence=_positive_int(value["fence"], "fence"),
|
|
909
|
+
capability_expires_at=expires_at,
|
|
910
|
+
callback_base_url=str(base_url).rstrip("/"),
|
|
911
|
+
allowed_callback_paths=tuple(paths),
|
|
912
|
+
heartbeat_interval_seconds=float(
|
|
913
|
+
_positive_int(value["heartbeat_interval_seconds"], "heartbeat_interval_seconds")
|
|
914
|
+
),
|
|
915
|
+
claim_poll_interval_seconds=float(
|
|
916
|
+
_positive_int(value["claim_poll_interval_seconds"], "claim_poll_interval_seconds")
|
|
917
|
+
),
|
|
918
|
+
idle_timeout_seconds=float(
|
|
919
|
+
_positive_int(value["idle_timeout_seconds"], "idle_timeout_seconds")
|
|
920
|
+
),
|
|
921
|
+
context=SessionContext(
|
|
922
|
+
dataset_id=_uuid_text(context["dataset_id"], "context.dataset_id"),
|
|
923
|
+
question_id=_uuid_text(context["question_id"], "context.question_id"),
|
|
924
|
+
table_id=(
|
|
925
|
+
_uuid_text(context["table_id"], "context.table_id")
|
|
926
|
+
if context.get("table_id") is not None
|
|
927
|
+
else None
|
|
928
|
+
),
|
|
929
|
+
source_ids=frozenset(
|
|
930
|
+
_uuid_text(item, "context.source_ids entry") for item in source_ids
|
|
931
|
+
),
|
|
932
|
+
),
|
|
933
|
+
)
|
|
934
|
+
except (KeyError, SessionWorkerError, HostedWorkerError, ValueError):
|
|
935
|
+
credential.close()
|
|
936
|
+
raise
|
|
937
|
+
return lease, credential
|
|
938
|
+
|
|
939
|
+
|
|
940
|
+
# ==========================================================================================
|
|
941
|
+
# Probe source access
|
|
942
|
+
# ==========================================================================================
|
|
943
|
+
|
|
944
|
+
|
|
945
|
+
@dataclass(frozen=True)
|
|
946
|
+
class ProbeSource:
|
|
947
|
+
"""The bytes a probe looks at, and where they lawfully came from."""
|
|
948
|
+
|
|
949
|
+
source_id: str
|
|
950
|
+
content: bytes
|
|
951
|
+
media_type: str
|
|
952
|
+
data_format: str
|
|
953
|
+
filename: str
|
|
954
|
+
truncated: bool
|
|
955
|
+
origin: str
|
|
956
|
+
|
|
957
|
+
|
|
958
|
+
class ProbeSourceReader(Protocol):
|
|
959
|
+
"""Where a probe's bytes come from. Injected, because which door is open is a deployment fact.
|
|
960
|
+
|
|
961
|
+
A session worker holds a *producer capability on a session run*, which is deliberately less
|
|
962
|
+
authority than a build attempt has, and that closes both of the older acquisition paths:
|
|
963
|
+
|
|
964
|
+
* ``producer/attempts/{id}/public-crawls`` (ADR 0015, worker side) requires
|
|
965
|
+
``table_recipe_id``, ``recipe_version``, ``recipe_digest``, ``activation_epoch`` and
|
|
966
|
+
``execution_scope_digest`` on the command, plus a frozen ``source_authority_digest`` from the
|
|
967
|
+
job's ``source_bindings``. A session run has none of the recipe-bound six -- Studio's own
|
|
968
|
+
design note says so, because probing is what a recipe gets written *from* -- so the command
|
|
969
|
+
cannot be constructed, let alone accepted.
|
|
970
|
+
* ``/v3/workspaces/{id}/public-acquisitions`` (ADR 0015, requester side) is
|
|
971
|
+
``userBearer``-authenticated. The worker holds no user bearer and must not be given one.
|
|
972
|
+
|
|
973
|
+
Studio built a third door for exactly this authority (Studio #99), and
|
|
974
|
+
:class:`StudioSessionAcquisitionReader` is the harness half of it. It is now the production
|
|
975
|
+
default. :class:`UnavailableProbeSourceReader` stays as the honest reader for a deployment
|
|
976
|
+
whose Studio predates that door: a probe that cannot read its source is a failed probe, not a
|
|
977
|
+
failed session, and the session stays live either way.
|
|
978
|
+
"""
|
|
979
|
+
|
|
980
|
+
def read(self, *, source_id: str, max_bytes: int) -> ProbeSource: ...
|
|
981
|
+
|
|
982
|
+
|
|
983
|
+
@dataclass(frozen=True)
|
|
984
|
+
class UnavailableProbeSourceReader:
|
|
985
|
+
"""Refuses every read, with the reason, and never guesses.
|
|
986
|
+
|
|
987
|
+
No longer the production default -- :class:`StudioSessionAcquisitionReader` is -- and kept
|
|
988
|
+
rather than deleted for the two cases where refusing is the correct answer: a Studio that
|
|
989
|
+
predates the session-acquisition door, and a test that wants the refusal path without a fake
|
|
990
|
+
for the whole acquisition protocol.
|
|
991
|
+
"""
|
|
992
|
+
|
|
993
|
+
reason: str = (
|
|
994
|
+
"no session-scoped source channel is wired: the producer public-crawl command requires "
|
|
995
|
+
"the recipe-bound coordinates a session run does not have, and the workspace "
|
|
996
|
+
"public-acquisition endpoint requires a user bearer this worker must not hold"
|
|
997
|
+
)
|
|
998
|
+
|
|
999
|
+
def read(self, *, source_id: str, max_bytes: int) -> ProbeSource:
|
|
1000
|
+
raise SessionWorkerError("PROBE_SOURCE_CHANNEL_UNAVAILABLE", self.reason)
|
|
1001
|
+
|
|
1002
|
+
|
|
1003
|
+
class StudioSessionAcquisitionReader:
|
|
1004
|
+
"""The production reader, against Studio's third acquisition door (Studio #99).
|
|
1005
|
+
|
|
1006
|
+
**Why this endpoint and not one of the two that already existed.** A session worker holds a
|
|
1007
|
+
producer capability on a *session* run. `producer/attempts/{id}/public-crawls` validates its
|
|
1008
|
+
command against a frozen recipe activation -- ``table_recipe_id``, ``recipe_version``,
|
|
1009
|
+
``recipe_digest``, ``activation_epoch``, ``execution_scope_digest`` and a
|
|
1010
|
+
``source_authority_digest`` from the run's source bindings -- and a session run carries none of
|
|
1011
|
+
the six, because probing is what a recipe gets written *from*. `/v3/workspaces/{id}/public-
|
|
1012
|
+
acquisitions` is ``userBearer``-authenticated and this worker must never hold a user bearer.
|
|
1013
|
+
`session-acquisitions` is the door built for exactly this authority, with two scopes of its own.
|
|
1014
|
+
|
|
1015
|
+
**The worker names a source and nothing else, and that is load-bearing rather than tidy.** The
|
|
1016
|
+
command admits no additional properties, so there is no field through which a URL, an adapter,
|
|
1017
|
+
a reader pin, a limit or an egress attestation could arrive. Studio resolves the locator from
|
|
1018
|
+
the source's one active ``public.https`` connector configuration -- a registration a *person*
|
|
1019
|
+
made -- applies its own reviewed egress policy, and clamps the registered limits to the session
|
|
1020
|
+
ceilings. A worker therefore cannot point the Clean room anywhere, and this class could not
|
|
1021
|
+
make it do so if it tried, because it has nothing to send.
|
|
1022
|
+
|
|
1023
|
+
**What is verified before anything is parsed.** The bundle is the same generation-pinned
|
|
1024
|
+
quarantine object the Builder collects, and the same obligation comes with it: the storage
|
|
1025
|
+
observation is not the evidence. `validate_result_download` re-derives the transfer authority
|
|
1026
|
+
and the pinned length and integrity, `verify_result_bytes` checks the bytes against them, and
|
|
1027
|
+
`parse_crawler_result` then re-derives the nested ``normalized_content_digest``,
|
|
1028
|
+
``normalized_size_bytes`` and ``result_digest`` from the document itself. Only after all three
|
|
1029
|
+
does a byte reach the sandbox.
|
|
1030
|
+
|
|
1031
|
+
**The wait is bounded and the loop is idempotent.** One create per source per attempt, under an
|
|
1032
|
+
Idempotency-Key derived from the source id, so a retried probe replays the first acquisition
|
|
1033
|
+
rather than spending a second admission slot -- the session's slots are charged to the person
|
|
1034
|
+
who opened it. Polling stops at `acquisition_timeout_seconds`; a probe is interactive and a
|
|
1035
|
+
crawl that has not settled by then is a timeout the client can be told about, not a thread the
|
|
1036
|
+
session should keep parked.
|
|
1037
|
+
"""
|
|
1038
|
+
|
|
1039
|
+
def __init__(
|
|
1040
|
+
self,
|
|
1041
|
+
*,
|
|
1042
|
+
authority: Any,
|
|
1043
|
+
models: Any,
|
|
1044
|
+
lease: SessionLease,
|
|
1045
|
+
transfer: Any,
|
|
1046
|
+
results_bucket: str | None = None,
|
|
1047
|
+
timeout_seconds: float = DEFAULT_ACQUISITION_TIMEOUT_SECONDS,
|
|
1048
|
+
poll_seconds: float = ACQUISITION_POLL_SECONDS,
|
|
1049
|
+
sleep: Callable[[float], None] = time.sleep,
|
|
1050
|
+
monotonic: Callable[[], float] = time.monotonic,
|
|
1051
|
+
now: Callable[[], datetime] = lambda: datetime.now(UTC),
|
|
1052
|
+
) -> None:
|
|
1053
|
+
self._authority = authority
|
|
1054
|
+
self._models = models
|
|
1055
|
+
self._lease = lease
|
|
1056
|
+
self._transfer = transfer
|
|
1057
|
+
self._results_bucket = results_bucket
|
|
1058
|
+
self._timeout = timeout_seconds
|
|
1059
|
+
self._poll = poll_seconds
|
|
1060
|
+
self._sleep = sleep
|
|
1061
|
+
self._monotonic = monotonic
|
|
1062
|
+
self._now = now
|
|
1063
|
+
# One entry per source, bumped only when that source's acquisition expired. It is what
|
|
1064
|
+
# makes the idempotency key stable for the ordinary case and re-openable for the one case
|
|
1065
|
+
# where replaying the old key would replay a corpse. See :meth:`read`.
|
|
1066
|
+
self._epochs: dict[str, int] = {}
|
|
1067
|
+
|
|
1068
|
+
# -- the read ----------------------------------------------------------------------------
|
|
1069
|
+
|
|
1070
|
+
def read(self, *, source_id: str, max_bytes: int) -> ProbeSource:
|
|
1071
|
+
try:
|
|
1072
|
+
session = self._settled(self._create(source_id), source_id)
|
|
1073
|
+
except SessionWorkerError as error:
|
|
1074
|
+
if error.code != ACQUISITION_EXPIRED_CODE:
|
|
1075
|
+
raise
|
|
1076
|
+
# An acquisition expires no later than the worker's attempt lease, and a warm session
|
|
1077
|
+
# outlives that comfortably -- the idle timeout alone is ten minutes. Without this, the
|
|
1078
|
+
# first source a long session reads becomes permanently unreadable the moment its
|
|
1079
|
+
# acquisition lapses, because the source-keyed idempotency key would keep replaying the
|
|
1080
|
+
# lapsed one. Studio releases the admission slot as it expires the acquisition, so
|
|
1081
|
+
# opening a fresh one under the next epoch is not a second spend. Exactly one retry:
|
|
1082
|
+
# a second expiry inside one probe is a clock problem, not a race.
|
|
1083
|
+
self._epochs[source_id] = self._epochs.get(source_id, 0) + 1
|
|
1084
|
+
session = self._settled(self._create(source_id), source_id)
|
|
1085
|
+
download = session.get("result_download")
|
|
1086
|
+
if not isinstance(download, Mapping):
|
|
1087
|
+
raise SessionWorkerError(
|
|
1088
|
+
"PROBE_SOURCE_RESULT_INVALID",
|
|
1089
|
+
"a ready session acquisition carried no result download",
|
|
1090
|
+
)
|
|
1091
|
+
raw = self._collect(download, session)
|
|
1092
|
+
return self._probe_source(raw, session, source_id, max_bytes)
|
|
1093
|
+
|
|
1094
|
+
# -- Studio ------------------------------------------------------------------------------
|
|
1095
|
+
|
|
1096
|
+
def _create(self, source_id: str) -> dict[str, Any]:
|
|
1097
|
+
command = self._models.CreateSessionScopedHostedAcquisitionCommand.from_dict(
|
|
1098
|
+
{
|
|
1099
|
+
"schema_version": START_SCHEMA_VERSION,
|
|
1100
|
+
"workspace_id": str(self._lease.workspace_id),
|
|
1101
|
+
"run_id": str(self._lease.run_id),
|
|
1102
|
+
"attempt_id": str(self._lease.attempt_id),
|
|
1103
|
+
"fence": self._lease.fence,
|
|
1104
|
+
"session_id": self._lease.session_id,
|
|
1105
|
+
"source_id": source_id,
|
|
1106
|
+
}
|
|
1107
|
+
)
|
|
1108
|
+
response = self._authority.create_session_acquisition.sync_detailed(
|
|
1109
|
+
attempt_id=self._lease.attempt_id,
|
|
1110
|
+
body=command,
|
|
1111
|
+
# KEYED ON THE SOURCE, deliberately, and not on a counter the way the claim poll
|
|
1112
|
+
# is. Two probes against one source inside one session must share the acquisition
|
|
1113
|
+
# they would otherwise both pay an admission slot for; a counter would turn a
|
|
1114
|
+
# follow-up question into a second crawl of the same bytes.
|
|
1115
|
+
idempotency_key=_idempotency(
|
|
1116
|
+
"session",
|
|
1117
|
+
self._lease.attempt_id,
|
|
1118
|
+
"acquire",
|
|
1119
|
+
f"{source_id}:{self._epochs.get(source_id, 0)}",
|
|
1120
|
+
),
|
|
1121
|
+
)
|
|
1122
|
+
try:
|
|
1123
|
+
return _response("createSessionAcquisition", response, 201).to_dict()
|
|
1124
|
+
except HostedWorkerError as error:
|
|
1125
|
+
if (
|
|
1126
|
+
error.code in PROBE_THROTTLE_WIRE_CODES
|
|
1127
|
+
and int(response.status_code) == HTTPStatus.TOO_MANY_REQUESTS
|
|
1128
|
+
):
|
|
1129
|
+
retry_after_seconds = _retry_after_seconds(getattr(response, "headers", {}))
|
|
1130
|
+
raise SessionWorkerError(
|
|
1131
|
+
error.code,
|
|
1132
|
+
error.detail,
|
|
1133
|
+
retry_after_seconds=retry_after_seconds,
|
|
1134
|
+
) from error
|
|
1135
|
+
raise
|
|
1136
|
+
|
|
1137
|
+
def _settled(self, session: Mapping[str, Any], source_id: str) -> dict[str, Any]:
|
|
1138
|
+
deadline = self._monotonic() + self._timeout
|
|
1139
|
+
try:
|
|
1140
|
+
crawler_session_id = UUID(str(session.get("crawler_session_id")))
|
|
1141
|
+
except ValueError as error:
|
|
1142
|
+
raise SessionWorkerError(
|
|
1143
|
+
"PROBE_SOURCE_ACQUISITION_MISBOUND",
|
|
1144
|
+
"session acquisition has no canonical crawler session id",
|
|
1145
|
+
) from error
|
|
1146
|
+
current = dict(session)
|
|
1147
|
+
while True:
|
|
1148
|
+
self._check(current, source_id)
|
|
1149
|
+
status = str(current.get("status", ""))
|
|
1150
|
+
if status == "result_ready":
|
|
1151
|
+
return current
|
|
1152
|
+
if status == "failed":
|
|
1153
|
+
raise SessionWorkerError(
|
|
1154
|
+
_acquisition_failure_code(current.get("failure_code")),
|
|
1155
|
+
"Studio session acquisition failed",
|
|
1156
|
+
)
|
|
1157
|
+
if self._monotonic() >= deadline:
|
|
1158
|
+
raise SessionWorkerError(
|
|
1159
|
+
"PROBE_SOURCE_ACQUISITION_TIMEOUT",
|
|
1160
|
+
f"session acquisition did not settle within {self._timeout:g}s",
|
|
1161
|
+
)
|
|
1162
|
+
self._sleep(self._poll)
|
|
1163
|
+
current = _response(
|
|
1164
|
+
"getSessionAcquisition",
|
|
1165
|
+
self._authority.get_session_acquisition.sync_detailed(
|
|
1166
|
+
attempt_id=self._lease.attempt_id,
|
|
1167
|
+
crawler_session_id=crawler_session_id,
|
|
1168
|
+
),
|
|
1169
|
+
200,
|
|
1170
|
+
).to_dict()
|
|
1171
|
+
|
|
1172
|
+
def _check(self, session: Mapping[str, Any], source_id: str) -> None:
|
|
1173
|
+
"""Refuse an acquisition that is not this attempt's, for this session, for this source.
|
|
1174
|
+
|
|
1175
|
+
Studio binds every one of these server-side, so this can only ever agree -- until the day
|
|
1176
|
+
it does not. Re-deriving them is what makes a mixed-up answer a refusal instead of bytes
|
|
1177
|
+
attributed to the wrong source, and it costs seven comparisons on a path that has just
|
|
1178
|
+
spent a network round trip.
|
|
1179
|
+
"""
|
|
1180
|
+
|
|
1181
|
+
if (
|
|
1182
|
+
str(session.get("workspace_id")) != str(self._lease.workspace_id)
|
|
1183
|
+
or str(session.get("run_id")) != str(self._lease.run_id)
|
|
1184
|
+
or str(session.get("attempt_id")) != str(self._lease.attempt_id)
|
|
1185
|
+
or str(session.get("research_session_id")) != str(self._lease.session_id)
|
|
1186
|
+
or session.get("source_id") != source_id
|
|
1187
|
+
or session.get("producer_fence") != self._lease.fence
|
|
1188
|
+
or session.get("adapter_id") != "public.https"
|
|
1189
|
+
):
|
|
1190
|
+
raise SessionWorkerError(
|
|
1191
|
+
"PROBE_SOURCE_ACQUISITION_MISBOUND",
|
|
1192
|
+
"session acquisition is not bound to this attempt, session and source",
|
|
1193
|
+
)
|
|
1194
|
+
|
|
1195
|
+
# -- the bundle --------------------------------------------------------------------------
|
|
1196
|
+
|
|
1197
|
+
def _collect(self, download: Mapping[str, Any], session: Mapping[str, Any]) -> bytes:
|
|
1198
|
+
declared = download.get("bundle_size_bytes")
|
|
1199
|
+
if type(declared) is not int or declared > MAX_PROBE_ACQUISITION_BUNDLE_BYTES:
|
|
1200
|
+
raise SessionWorkerError(
|
|
1201
|
+
"PROBE_SOURCE_TOO_LARGE",
|
|
1202
|
+
"the acquired bundle is larger than a probe may hold",
|
|
1203
|
+
)
|
|
1204
|
+
try:
|
|
1205
|
+
validated = validate_result_download(
|
|
1206
|
+
download,
|
|
1207
|
+
base_url=self._lease.callback_base_url,
|
|
1208
|
+
crawler_session_id=str(session["crawler_session_id"]),
|
|
1209
|
+
request_digest=str(session["request_digest"]),
|
|
1210
|
+
source_id=str(session["source_id"]),
|
|
1211
|
+
source_authority_digest=str(session["source_authority_digest"]),
|
|
1212
|
+
now=self._now(),
|
|
1213
|
+
expected_gcs_bucket=self._results_bucket,
|
|
1214
|
+
)
|
|
1215
|
+
except ResultDownloadError as error:
|
|
1216
|
+
raise SessionWorkerError("PROBE_SOURCE_DOWNLOAD_INVALID", str(error)) from error
|
|
1217
|
+
content = bytearray()
|
|
1218
|
+
with self._transfer.stream("GET", validated.url, headers={}) as response:
|
|
1219
|
+
if response.status_code != 200:
|
|
1220
|
+
raise SessionWorkerError(
|
|
1221
|
+
"PROBE_SOURCE_DOWNLOAD_FAILED",
|
|
1222
|
+
f"quarantine download returned HTTP {response.status_code}",
|
|
1223
|
+
)
|
|
1224
|
+
for chunk in response.iter_bytes():
|
|
1225
|
+
content.extend(chunk)
|
|
1226
|
+
if len(content) > validated.size:
|
|
1227
|
+
raise SessionWorkerError(
|
|
1228
|
+
"PROBE_SOURCE_DOWNLOAD_FAILED",
|
|
1229
|
+
"quarantine bytes exceed the generation-pinned size",
|
|
1230
|
+
)
|
|
1231
|
+
raw = bytes(content)
|
|
1232
|
+
try:
|
|
1233
|
+
verify_result_bytes(validated, raw)
|
|
1234
|
+
except ResultDownloadError as error:
|
|
1235
|
+
raise SessionWorkerError("PROBE_SOURCE_DOWNLOAD_FAILED", str(error)) from error
|
|
1236
|
+
return raw
|
|
1237
|
+
|
|
1238
|
+
def _probe_source(
|
|
1239
|
+
self, raw: bytes, session: Mapping[str, Any], source_id: str, max_bytes: int
|
|
1240
|
+
) -> ProbeSource:
|
|
1241
|
+
document = _strict_json(raw, "session acquisition result")
|
|
1242
|
+
if not isinstance(document, dict) or canonical_json_bytes(document) != raw:
|
|
1243
|
+
raise SessionWorkerError(
|
|
1244
|
+
"PROBE_SOURCE_RESULT_INVALID",
|
|
1245
|
+
"the acquired result is not exact canonical JSON",
|
|
1246
|
+
)
|
|
1247
|
+
try:
|
|
1248
|
+
result = parse_crawler_result(document)
|
|
1249
|
+
except HostedCrawlerProtocolError as error:
|
|
1250
|
+
raise SessionWorkerError("PROBE_SOURCE_RESULT_INVALID", error.detail) from error
|
|
1251
|
+
receipt = result.acquisition_receipt
|
|
1252
|
+
if (
|
|
1253
|
+
result.source_id != source_id
|
|
1254
|
+
or result.request_digest != session.get("request_digest")
|
|
1255
|
+
or result.adapter_id != "public.https"
|
|
1256
|
+
or receipt.get("source_id") != source_id
|
|
1257
|
+
or receipt.get("normalized_content_sha256") != sha256_bytes(result.normalized_content)
|
|
1258
|
+
or receipt.get("normalized_size_bytes") != len(result.normalized_content)
|
|
1259
|
+
or receipt.get("egress_policy_attestation") != session.get("egress_policy_attestation")
|
|
1260
|
+
):
|
|
1261
|
+
raise SessionWorkerError(
|
|
1262
|
+
"PROBE_SOURCE_RESULT_INVALID",
|
|
1263
|
+
"the acquired result is not bound to the acquisition that produced it",
|
|
1264
|
+
)
|
|
1265
|
+
if len(result.normalized_content) > max_bytes:
|
|
1266
|
+
raise SessionWorkerError(
|
|
1267
|
+
"PROBE_SOURCE_TOO_LARGE",
|
|
1268
|
+
"the acquired source is larger than a probe may hold",
|
|
1269
|
+
)
|
|
1270
|
+
media_type = receipt.get("normalized_media_type")
|
|
1271
|
+
data_format = receipt.get("normalized_data_format")
|
|
1272
|
+
filename = receipt.get("normalized_filename")
|
|
1273
|
+
if (
|
|
1274
|
+
not isinstance(media_type, str)
|
|
1275
|
+
or not isinstance(data_format, str)
|
|
1276
|
+
or not isinstance(filename, str)
|
|
1277
|
+
):
|
|
1278
|
+
raise SessionWorkerError(
|
|
1279
|
+
"PROBE_SOURCE_RESULT_INVALID",
|
|
1280
|
+
"the acquisition receipt does not name its normalized shape",
|
|
1281
|
+
)
|
|
1282
|
+
return ProbeSource(
|
|
1283
|
+
source_id=source_id,
|
|
1284
|
+
content=result.normalized_content,
|
|
1285
|
+
media_type=media_type,
|
|
1286
|
+
data_format=data_format,
|
|
1287
|
+
filename=filename,
|
|
1288
|
+
# The bytes are whole or they are refused. Studio clamps the source to the session
|
|
1289
|
+
# ceiling before it crawls, and a truncated table would make every count a probe
|
|
1290
|
+
# reports quietly wrong, so there is no partial answer to flag here.
|
|
1291
|
+
truncated=False,
|
|
1292
|
+
origin="studio.session-acquisition",
|
|
1293
|
+
)
|
|
1294
|
+
|
|
1295
|
+
|
|
1296
|
+
def _acquisition_failure_code(value: Any) -> str:
|
|
1297
|
+
"""Render Studio's ``failure_code`` as a code this worker will settle a probe under.
|
|
1298
|
+
|
|
1299
|
+
Studio's is free-form on the wire; anything that is not a code-shaped token becomes one stable
|
|
1300
|
+
value rather than reaching the probe record, and the timeline, as somebody's prose.
|
|
1301
|
+
"""
|
|
1302
|
+
|
|
1303
|
+
if isinstance(value, str) and re.fullmatch(r"[A-Z][A-Z0-9_]{2,127}", value):
|
|
1304
|
+
return value
|
|
1305
|
+
return "PROBE_SOURCE_ACQUISITION_FAILED"
|
|
1306
|
+
|
|
1307
|
+
|
|
1308
|
+
def install_hosted_catalogue() -> Mapping[str, Any] | None:
|
|
1309
|
+
"""Put this deployment's pinned sealed catalogue on the worker, when one is pinned.
|
|
1310
|
+
|
|
1311
|
+
The same hook the Builder calls, at start rather than at first query, so a worker that needs
|
|
1312
|
+
the catalogue discovers it cannot have one while it can still fail cleanly. ``None`` when no
|
|
1313
|
+
generation is pinned, which is an ordinary deployment and not a fault: probes then run against
|
|
1314
|
+
whatever the source reader provides.
|
|
1315
|
+
|
|
1316
|
+
Imported inside the function rather than at module scope. The import reaches the channel
|
|
1317
|
+
client and its digest verification, and a session that never touches the catalogue should not
|
|
1318
|
+
pay for loading them.
|
|
1319
|
+
"""
|
|
1320
|
+
|
|
1321
|
+
from mostlyright.data_harness.sources.catalog.hosted_catalog import (
|
|
1322
|
+
ensure_hosted_catalog,
|
|
1323
|
+
)
|
|
1324
|
+
|
|
1325
|
+
receipt = ensure_hosted_catalog()
|
|
1326
|
+
return None if receipt is None else dict(receipt)
|
|
1327
|
+
|
|
1328
|
+
|
|
1329
|
+
# ==========================================================================================
|
|
1330
|
+
# Probe execution
|
|
1331
|
+
# ==========================================================================================
|
|
1332
|
+
|
|
1333
|
+
|
|
1334
|
+
class ProbeExecutor:
|
|
1335
|
+
"""Run one probe inside the sandbox and answer with a vocabulary-valid result body.
|
|
1336
|
+
|
|
1337
|
+
Every probe is a *no-network* sandbox operation (``parse``), which is the only confinement the
|
|
1338
|
+
warm pool serves and therefore the only one that can meet the follow-up-probe budget. Bytes
|
|
1339
|
+
reach the executor already acquired, through the injected :class:`ProbeSourceReader`; nothing
|
|
1340
|
+
here retrieves anything, so a probe cannot cause an outbound request at all.
|
|
1341
|
+
|
|
1342
|
+
**The read is before the dispatch because all four kinds need it, not by accident.** When the
|
|
1343
|
+
read refuses -- against a Studio with no session-acquisition door, or for a source with no
|
|
1344
|
+
registration -- every kind refuses with it, and the obvious question is whether some kind could
|
|
1345
|
+
have been answered from metadata instead and let through. The answer is no, for each of them:
|
|
1346
|
+
|
|
1347
|
+
* ``sample_rows`` and ``profile_columns`` answer with cell values and per-column statistics
|
|
1348
|
+
computed over them.
|
|
1349
|
+
* ``evaluate_expression`` -- the one that looks like it might, since ``count`` with no
|
|
1350
|
+
predicate is just a number -- filters *rows* and aggregates *values*. Even the degenerate
|
|
1351
|
+
case would have to answer from ``declared_row_count``, which is what a publisher SAYS, and
|
|
1352
|
+
answering an observation with a claim is the substitution this whole lane exists to avoid.
|
|
1353
|
+
* ``source_inspect`` is the closest call. The sealed catalogue does carry ``declared_columns``,
|
|
1354
|
+
but they are provider-declared names with no types; the answer this kind gives is a
|
|
1355
|
+
``schema_digest`` from the parse plus an ``inferred_type`` per column, both derived from the
|
|
1356
|
+
scanned cells, and neither is reachable without them.
|
|
1357
|
+
|
|
1358
|
+
So a read refusal is per-read and not per-kind, and there is no subset to exempt. What a caller
|
|
1359
|
+
gets instead is a *settled* probe carrying the refusal's own code -- reported to Studio so the
|
|
1360
|
+
probe record terminates, and narrated on the run stream as an ``outcome="refused"`` completion
|
|
1361
|
+
-- rather than a probe that starts and never ends.
|
|
1362
|
+
"""
|
|
1363
|
+
|
|
1364
|
+
def __init__(
|
|
1365
|
+
self,
|
|
1366
|
+
*,
|
|
1367
|
+
sandbox: Any,
|
|
1368
|
+
reader: ProbeSourceReader,
|
|
1369
|
+
context: SessionContext,
|
|
1370
|
+
parse_limits: ParseLimits | None = None,
|
|
1371
|
+
max_source_bytes: int = 16 * 1024 * 1024,
|
|
1372
|
+
) -> None:
|
|
1373
|
+
self._sandbox = sandbox
|
|
1374
|
+
self._reader = reader
|
|
1375
|
+
self._context = context
|
|
1376
|
+
self._limits = parse_limits or ParseLimits()
|
|
1377
|
+
self._max_source_bytes = max_source_bytes
|
|
1378
|
+
|
|
1379
|
+
def execute(self, *, probe_id: str, probe_kind: str, request: Any) -> dict[str, Any]:
|
|
1380
|
+
"""Answer one probe.
|
|
1381
|
+
|
|
1382
|
+
Raises :class:`SessionWorkerError`, or
|
|
1383
|
+
:class:`~mostlyright.data_harness.session_probes.ProbeVocabularyError`, or whatever the
|
|
1384
|
+
Clean room raises -- every one of which carries a ``code`` that
|
|
1385
|
+
:func:`_probe_failure_code` settles the probe under.
|
|
1386
|
+
"""
|
|
1387
|
+
|
|
1388
|
+
normalized = session_probes.validate_probe_request(probe_kind, request)
|
|
1389
|
+
source_id = normalized["source_id"]
|
|
1390
|
+
# The session's own context is the probe authority. A probe naming a source the session was
|
|
1391
|
+
# not opened against is refused here, before any read: the capability says which session
|
|
1392
|
+
# this is, and the session says which sources it may look at. Studio does not re-check this
|
|
1393
|
+
# -- the probe body is opaque to it -- so if this check is not here, it is nowhere.
|
|
1394
|
+
if self._context.source_ids and source_id not in self._context.source_ids:
|
|
1395
|
+
raise SessionWorkerError(
|
|
1396
|
+
"PROBE_SOURCE_OUT_OF_SCOPE",
|
|
1397
|
+
"probe names a source this session was not opened against",
|
|
1398
|
+
)
|
|
1399
|
+
source = self._reader.read(source_id=source_id, max_bytes=self._max_source_bytes)
|
|
1400
|
+
if source.source_id != source_id:
|
|
1401
|
+
raise SessionWorkerError(
|
|
1402
|
+
"PROBE_SOURCE_MISMATCH", "source reader answered for a different source"
|
|
1403
|
+
)
|
|
1404
|
+
result = self._sandbox.parse(
|
|
1405
|
+
request_id=_sandbox_request_id(probe_id),
|
|
1406
|
+
content=source.content,
|
|
1407
|
+
data_format=source.data_format,
|
|
1408
|
+
media_type=source.media_type,
|
|
1409
|
+
filename=source.filename,
|
|
1410
|
+
limits=self._limits,
|
|
1411
|
+
)
|
|
1412
|
+
parsed = getattr(result, "parsed", None)
|
|
1413
|
+
if parsed is None:
|
|
1414
|
+
raise SessionWorkerError("PROBE_PARSE_EMPTY", "the Clean room returned no table")
|
|
1415
|
+
envelope = {
|
|
1416
|
+
"schema_version": session_probes.PROBE_SCHEMA_VERSION,
|
|
1417
|
+
"probe_kind": probe_kind,
|
|
1418
|
+
"source_id": source_id,
|
|
1419
|
+
"content_sha256": hashlib.sha256(source.content).hexdigest(),
|
|
1420
|
+
"content_bytes": len(source.content),
|
|
1421
|
+
"data_format": parsed.data_format,
|
|
1422
|
+
"sandbox_policy_digest": result.policy_digest,
|
|
1423
|
+
"truncated": bool(source.truncated),
|
|
1424
|
+
}
|
|
1425
|
+
body = self._answer(probe_kind, normalized, parsed, envelope)
|
|
1426
|
+
return session_probes.validate_probe_result(probe_kind, body)
|
|
1427
|
+
|
|
1428
|
+
# -- per-kind answers -------------------------------------------------------------------
|
|
1429
|
+
|
|
1430
|
+
def _answer(
|
|
1431
|
+
self,
|
|
1432
|
+
probe_kind: str,
|
|
1433
|
+
request: Mapping[str, Any],
|
|
1434
|
+
parsed: Any,
|
|
1435
|
+
envelope: dict[str, Any],
|
|
1436
|
+
) -> dict[str, Any]:
|
|
1437
|
+
columns: tuple[str, ...] = tuple(parsed.columns)
|
|
1438
|
+
rows: tuple[tuple[Any, ...], ...] = tuple(parsed.rows)
|
|
1439
|
+
if probe_kind == "sample_rows":
|
|
1440
|
+
return self._sample_rows(request, columns, rows, envelope)
|
|
1441
|
+
scanned = rows[: request["scan_rows"]]
|
|
1442
|
+
envelope["rows_scanned"] = len(scanned)
|
|
1443
|
+
if probe_kind == "source_inspect":
|
|
1444
|
+
return self._source_inspect(columns, scanned, parsed, envelope)
|
|
1445
|
+
if probe_kind == "profile_columns":
|
|
1446
|
+
return self._profile_columns(request, columns, scanned, envelope)
|
|
1447
|
+
return self._evaluate_expression(request, columns, scanned, envelope)
|
|
1448
|
+
|
|
1449
|
+
def _source_inspect(
|
|
1450
|
+
self,
|
|
1451
|
+
columns: Sequence[str],
|
|
1452
|
+
scanned: Sequence[Sequence[Any]],
|
|
1453
|
+
parsed: Any,
|
|
1454
|
+
envelope: dict[str, Any],
|
|
1455
|
+
) -> dict[str, Any]:
|
|
1456
|
+
# Deliberately no cell values. `source_inspect` answers "what shape is this", and a client
|
|
1457
|
+
# that wants values asks `sample_rows` -- which is a separate probe, separately visible on
|
|
1458
|
+
# the session's event stream, so "who looked at the data" stays answerable.
|
|
1459
|
+
return {
|
|
1460
|
+
**envelope,
|
|
1461
|
+
"schema_digest": parsed.schema_digest,
|
|
1462
|
+
"columns": [
|
|
1463
|
+
{
|
|
1464
|
+
"name": name,
|
|
1465
|
+
"inferred_type": session_probes.inferred_type([row[index] for row in scanned]),
|
|
1466
|
+
}
|
|
1467
|
+
for index, name in enumerate(columns[: session_probes.MAX_PROFILE_COLUMNS])
|
|
1468
|
+
],
|
|
1469
|
+
}
|
|
1470
|
+
|
|
1471
|
+
def _sample_rows(
|
|
1472
|
+
self,
|
|
1473
|
+
request: Mapping[str, Any],
|
|
1474
|
+
columns: Sequence[str],
|
|
1475
|
+
rows: Sequence[Sequence[Any]],
|
|
1476
|
+
envelope: dict[str, Any],
|
|
1477
|
+
) -> dict[str, Any]:
|
|
1478
|
+
selected = self._selected_columns(request, columns)
|
|
1479
|
+
indices = [columns.index(name) for name in selected]
|
|
1480
|
+
window = rows[request["offset"] : request["offset"] + request["limit"]]
|
|
1481
|
+
envelope["rows_scanned"] = request["offset"] + len(window)
|
|
1482
|
+
body = {
|
|
1483
|
+
**envelope,
|
|
1484
|
+
"columns": list(selected),
|
|
1485
|
+
"rows": [
|
|
1486
|
+
[session_probes.preview_cell(row[index]) for index in indices] for row in window
|
|
1487
|
+
],
|
|
1488
|
+
}
|
|
1489
|
+
# A preview that would not fit is trimmed from the end rather than refused: an interactive
|
|
1490
|
+
# caller asking for 200 wide rows should get the first N it can have, and `truncated` says
|
|
1491
|
+
# so. Refusing would make the useful answer unreachable without a second round trip.
|
|
1492
|
+
while (
|
|
1493
|
+
len(canonical_json_bytes(body)) > session_probes.PROBE_RESULT_MAX_BYTES and body["rows"]
|
|
1494
|
+
):
|
|
1495
|
+
body["rows"].pop()
|
|
1496
|
+
body["truncated"] = True
|
|
1497
|
+
return body
|
|
1498
|
+
|
|
1499
|
+
def _profile_columns(
|
|
1500
|
+
self,
|
|
1501
|
+
request: Mapping[str, Any],
|
|
1502
|
+
columns: Sequence[str],
|
|
1503
|
+
scanned: Sequence[Sequence[Any]],
|
|
1504
|
+
envelope: dict[str, Any],
|
|
1505
|
+
) -> dict[str, Any]:
|
|
1506
|
+
selected = self._selected_columns(request, columns)
|
|
1507
|
+
profiles: list[dict[str, Any]] = []
|
|
1508
|
+
for name in selected:
|
|
1509
|
+
index = columns.index(name)
|
|
1510
|
+
values = [row[index] for row in scanned]
|
|
1511
|
+
present = [value for value in values if value is not None]
|
|
1512
|
+
profile: dict[str, Any] = {
|
|
1513
|
+
"name": name,
|
|
1514
|
+
"inferred_type": session_probes.inferred_type(values),
|
|
1515
|
+
"non_null": len(present),
|
|
1516
|
+
"nulls": len(values) - len(present),
|
|
1517
|
+
# Exact over the scanned rows, and `rows_scanned` is in the envelope beside it, so
|
|
1518
|
+
# the number always arrives with its denominator.
|
|
1519
|
+
"distinct": len({session_probes.preview_cell(value) for value in present}),
|
|
1520
|
+
}
|
|
1521
|
+
ordered = self._orderable(present)
|
|
1522
|
+
if ordered:
|
|
1523
|
+
profile["minimum"] = session_probes.preview_cell(ordered[0])
|
|
1524
|
+
profile["maximum"] = session_probes.preview_cell(ordered[-1])
|
|
1525
|
+
profiles.append(profile)
|
|
1526
|
+
return {**envelope, "columns": profiles}
|
|
1527
|
+
|
|
1528
|
+
def _evaluate_expression(
|
|
1529
|
+
self,
|
|
1530
|
+
request: Mapping[str, Any],
|
|
1531
|
+
columns: Sequence[str],
|
|
1532
|
+
scanned: Sequence[Sequence[Any]],
|
|
1533
|
+
envelope: dict[str, Any],
|
|
1534
|
+
) -> dict[str, Any]:
|
|
1535
|
+
expression = request["expression"]
|
|
1536
|
+
predicate = expression.get("where")
|
|
1537
|
+
aggregate = expression["aggregate"]
|
|
1538
|
+
column = expression.get("column")
|
|
1539
|
+
index = columns.index(column) if column in columns else None
|
|
1540
|
+
matched = 0
|
|
1541
|
+
values: list[Any] = []
|
|
1542
|
+
for row in scanned:
|
|
1543
|
+
record = dict(zip(columns, row, strict=False))
|
|
1544
|
+
if predicate is not None and not session_probes.evaluate_predicate(predicate, record):
|
|
1545
|
+
continue
|
|
1546
|
+
matched += 1
|
|
1547
|
+
if index is not None and row[index] is not None:
|
|
1548
|
+
values.append(row[index])
|
|
1549
|
+
body: dict[str, Any] = {
|
|
1550
|
+
**envelope,
|
|
1551
|
+
"aggregate": aggregate,
|
|
1552
|
+
"value": self._aggregate(aggregate, matched, values),
|
|
1553
|
+
"rows_matched": matched,
|
|
1554
|
+
}
|
|
1555
|
+
if column is not None:
|
|
1556
|
+
body["column"] = column
|
|
1557
|
+
return body
|
|
1558
|
+
|
|
1559
|
+
@staticmethod
|
|
1560
|
+
def _aggregate(aggregate: str, matched: int, values: Sequence[Any]) -> Any:
|
|
1561
|
+
if aggregate == "count":
|
|
1562
|
+
return matched
|
|
1563
|
+
if not values:
|
|
1564
|
+
# No value is `null`, never zero. A sum over nothing that reported 0 would be
|
|
1565
|
+
# indistinguishable from a sum over zeros, and an exploratory probe is exactly where
|
|
1566
|
+
# that distinction matters.
|
|
1567
|
+
return None
|
|
1568
|
+
numeric = [value for value in values if type(value) in {int, float}]
|
|
1569
|
+
if aggregate in {"min", "max"}:
|
|
1570
|
+
ordered = ProbeExecutor._orderable(values)
|
|
1571
|
+
if not ordered:
|
|
1572
|
+
return None
|
|
1573
|
+
return session_probes.preview_cell(ordered[0 if aggregate == "min" else -1])
|
|
1574
|
+
if len(numeric) != len(values):
|
|
1575
|
+
raise SessionWorkerError(
|
|
1576
|
+
"PROBE_AGGREGATE_NOT_NUMERIC",
|
|
1577
|
+
f"{aggregate} needs a numeric column; the scanned values are not all numbers",
|
|
1578
|
+
)
|
|
1579
|
+
total = sum(numeric)
|
|
1580
|
+
return session_probes.preview_cell(total if aggregate == "sum" else total / len(numeric))
|
|
1581
|
+
|
|
1582
|
+
@staticmethod
|
|
1583
|
+
def _orderable(values: Sequence[Any]) -> list[Any]:
|
|
1584
|
+
"""Sort a column's present values, or answer empty when the column is not orderable.
|
|
1585
|
+
|
|
1586
|
+
Python 3 raises on ``1 < "a"``; a mixed column therefore has no extremes and says so,
|
|
1587
|
+
rather than the probe failing on data the caller cannot be expected to know the shape of.
|
|
1588
|
+
"""
|
|
1589
|
+
|
|
1590
|
+
try:
|
|
1591
|
+
return sorted(values)
|
|
1592
|
+
except TypeError:
|
|
1593
|
+
return []
|
|
1594
|
+
|
|
1595
|
+
def _selected_columns(
|
|
1596
|
+
self, request: Mapping[str, Any], columns: Sequence[str]
|
|
1597
|
+
) -> tuple[str, ...]:
|
|
1598
|
+
requested = request.get("columns")
|
|
1599
|
+
if requested is None:
|
|
1600
|
+
return tuple(columns[: session_probes.MAX_PROFILE_COLUMNS])
|
|
1601
|
+
unknown = [name for name in requested if name not in columns]
|
|
1602
|
+
if unknown:
|
|
1603
|
+
raise SessionWorkerError(
|
|
1604
|
+
"PROBE_COLUMN_UNKNOWN", f"source has no column named {unknown[0]!r}"
|
|
1605
|
+
)
|
|
1606
|
+
return tuple(requested)
|
|
1607
|
+
|
|
1608
|
+
|
|
1609
|
+
def _probe_failure_code(error: BaseException) -> str:
|
|
1610
|
+
"""Render one probe failure as the stable code it should settle under.
|
|
1611
|
+
|
|
1612
|
+
A ``code`` attribute is the convention every coded failure in this package already follows --
|
|
1613
|
+
:class:`SessionWorkerError`, :class:`~mostlyright.data_harness.session_probes.
|
|
1614
|
+
ProbeVocabularyError`, :class:`~mostlyright.data_harness.hosted_worker.HostedWorkerError` and
|
|
1615
|
+
the Clean room's ``AcquisitionSecurityError`` all carry one -- so this reads the convention
|
|
1616
|
+
rather than a list of classes. Reading the class list instead is what lost ``SANDBOX_TIMEOUT``:
|
|
1617
|
+
the sandbox's error is declared in ``acquisition.url_policy`` and was in nobody's tuple.
|
|
1618
|
+
|
|
1619
|
+
The value is checked against the progress value domain before it is believed. A code reaches a
|
|
1620
|
+
timeline as a fact, and a fact that is prose is a record the vocabulary drops whole -- which
|
|
1621
|
+
would put the pair invariant back exactly where this change found it.
|
|
1622
|
+
"""
|
|
1623
|
+
|
|
1624
|
+
code = getattr(error, "code", None)
|
|
1625
|
+
if isinstance(code, str) and progress_events.is_progress_fact_value(code):
|
|
1626
|
+
return code
|
|
1627
|
+
return "PROBE_FAILED"
|
|
1628
|
+
|
|
1629
|
+
|
|
1630
|
+
def _retry_after_seconds(headers: Any) -> int | None:
|
|
1631
|
+
"""Accept Studio's bounded delta-seconds retry window, never prose or a guessed delay."""
|
|
1632
|
+
|
|
1633
|
+
if not isinstance(headers, Mapping):
|
|
1634
|
+
return None
|
|
1635
|
+
value = next(
|
|
1636
|
+
(
|
|
1637
|
+
candidate
|
|
1638
|
+
for key, candidate in headers.items()
|
|
1639
|
+
if isinstance(key, str) and key.lower() == "retry-after"
|
|
1640
|
+
),
|
|
1641
|
+
None,
|
|
1642
|
+
)
|
|
1643
|
+
if not isinstance(value, str) or not re.fullmatch(r"[1-9][0-9]{0,3}", value):
|
|
1644
|
+
return None
|
|
1645
|
+
seconds = int(value)
|
|
1646
|
+
return seconds if seconds <= MAX_RETRY_AFTER_SECONDS else None
|
|
1647
|
+
|
|
1648
|
+
|
|
1649
|
+
def _probe_display_failure(error: BaseException, wire_code: str) -> tuple[str, int]:
|
|
1650
|
+
"""Return a person-facing code and its server-stated wait while retaining wire identity."""
|
|
1651
|
+
|
|
1652
|
+
if wire_code not in PROBE_THROTTLE_WIRE_CODES:
|
|
1653
|
+
return wire_code, 0
|
|
1654
|
+
retry_after_seconds = getattr(error, "retry_after_seconds", None)
|
|
1655
|
+
if type(retry_after_seconds) is int and 1 <= retry_after_seconds <= MAX_RETRY_AFTER_SECONDS:
|
|
1656
|
+
return SESSION_ACQUISITION_THROTTLED, retry_after_seconds
|
|
1657
|
+
return SESSION_ACQUISITION_THROTTLED, 0
|
|
1658
|
+
|
|
1659
|
+
|
|
1660
|
+
def _probe_progress_outcome(failure_code: str | None) -> str:
|
|
1661
|
+
"""Classify a settled probe for the unsealed timeline. Total over ``PROBE_OUTCOMES``.
|
|
1662
|
+
|
|
1663
|
+
Keyed on the code rather than on the exception, so the classification is a property of what the
|
|
1664
|
+
probe settled AS -- the same thing the client reads off the probe resource -- and not of how
|
|
1665
|
+
the worker happened to learn it.
|
|
1666
|
+
"""
|
|
1667
|
+
|
|
1668
|
+
if failure_code is None:
|
|
1669
|
+
return "succeeded"
|
|
1670
|
+
if failure_code in PROBE_TIMEOUT_CODES:
|
|
1671
|
+
return "timed_out"
|
|
1672
|
+
if failure_code in PROBE_THROTTLE_CODES:
|
|
1673
|
+
return "throttled"
|
|
1674
|
+
if failure_code in PROBE_REFUSAL_CODES:
|
|
1675
|
+
return "refused"
|
|
1676
|
+
return "failed"
|
|
1677
|
+
|
|
1678
|
+
|
|
1679
|
+
# ==========================================================================================
|
|
1680
|
+
# The Studio authority
|
|
1681
|
+
# ==========================================================================================
|
|
1682
|
+
|
|
1683
|
+
|
|
1684
|
+
def _sandbox_request_id(probe_id: str) -> str:
|
|
1685
|
+
"""Render a probe id as a sandbox request id.
|
|
1686
|
+
|
|
1687
|
+
``CrawlerSandbox`` requires a lowercase bounded identifier that *starts with a letter*, and a
|
|
1688
|
+
UUID beginning with a digit does not -- which is roughly half of them. Passing the probe id
|
|
1689
|
+
straight through therefore refused every other probe with ``SANDBOX_IDENTIFIER``, and the
|
|
1690
|
+
prefix is what fixes it rather than a relaxation on the sandbox side. Anything outside the
|
|
1691
|
+
admitted alphabet is replaced rather than dropped, so two different probe ids cannot collapse
|
|
1692
|
+
into one request id.
|
|
1693
|
+
"""
|
|
1694
|
+
|
|
1695
|
+
safe = "".join(
|
|
1696
|
+
character if character.islower() or character.isdigit() or character in "._-" else "-"
|
|
1697
|
+
for character in probe_id.lower()
|
|
1698
|
+
)
|
|
1699
|
+
return f"probe-{safe}"[:128]
|
|
1700
|
+
|
|
1701
|
+
|
|
1702
|
+
class _SessionAuthority(_ProducerAuthority):
|
|
1703
|
+
"""The producer surface plus the two verbs a session adds, and the lease that precedes both.
|
|
1704
|
+
|
|
1705
|
+
Subclassed here rather than widened in ``hosted_worker`` so that the producer and verifier
|
|
1706
|
+
façades keep exactly the operation sets they had: a session worker gaining a verb must not
|
|
1707
|
+
silently widen what a build worker's façade will dispatch.
|
|
1708
|
+
"""
|
|
1709
|
+
|
|
1710
|
+
_operations: ClassVar[dict[str, str]] = {
|
|
1711
|
+
**_ProducerAuthority._operations,
|
|
1712
|
+
"lease_session": "sessions.lease_research_session",
|
|
1713
|
+
"claim_probe": "producer.claim_research_probe",
|
|
1714
|
+
"report_probe_result": "producer.report_research_probe_result",
|
|
1715
|
+
"create_session_acquisition": "producer.create_session_acquisition",
|
|
1716
|
+
"get_session_acquisition": "producer.get_session_acquisition",
|
|
1717
|
+
}
|
|
1718
|
+
|
|
1719
|
+
|
|
1720
|
+
def _lease_session(
|
|
1721
|
+
*,
|
|
1722
|
+
config: SessionWorkerConfig,
|
|
1723
|
+
invocation: StartInvocation,
|
|
1724
|
+
nonce: bytearray,
|
|
1725
|
+
transport: BootstrapTransport,
|
|
1726
|
+
) -> tuple[SessionLease, CredentialLease]:
|
|
1727
|
+
"""Present ambient workload identity plus the nonce, and take the one lease on this session.
|
|
1728
|
+
|
|
1729
|
+
Two credentials are in play for exactly one call each and both are zeroed. The identity token
|
|
1730
|
+
proves *which fleet* is calling and is destroyed the moment the exchange returns. The nonce
|
|
1731
|
+
proves *which invocation*, and is zeroed by the caller as soon as this returns -- it is a
|
|
1732
|
+
one-time secret whose whole purpose is to stop a reviewed worker leasing a session it was not
|
|
1733
|
+
sent to, and a copy that outlives its one use is a copy that can be stolen.
|
|
1734
|
+
"""
|
|
1735
|
+
|
|
1736
|
+
models, _, _, _ = _load_generated()
|
|
1737
|
+
identity = _metadata_identity(transport, config.lease_audience)
|
|
1738
|
+
try:
|
|
1739
|
+
client = _authenticated_client(config.studio_origin, identity.decode("ascii"))
|
|
1740
|
+
finally:
|
|
1741
|
+
identity[:] = b"\x00" * len(identity)
|
|
1742
|
+
try:
|
|
1743
|
+
authority = _SessionAuthority(client)
|
|
1744
|
+
command = models.ContractSessionLeaseCommand.from_dict(
|
|
1745
|
+
{
|
|
1746
|
+
"schema_version": START_SCHEMA_VERSION,
|
|
1747
|
+
"session_id": invocation.session_id,
|
|
1748
|
+
"lease_nonce": nonce.decode("ascii"),
|
|
1749
|
+
}
|
|
1750
|
+
)
|
|
1751
|
+
parsed = _response(
|
|
1752
|
+
"leaseResearchSession",
|
|
1753
|
+
authority.lease_session.sync_detailed(
|
|
1754
|
+
session_id=UUID(invocation.session_id), body=command
|
|
1755
|
+
),
|
|
1756
|
+
200,
|
|
1757
|
+
)
|
|
1758
|
+
wire = parsed.to_dict()
|
|
1759
|
+
finally:
|
|
1760
|
+
# The identity-bearing client is discarded whatever happens: everything after the lease
|
|
1761
|
+
# authenticates with the fenced capability, and an idle client holding a workload-identity
|
|
1762
|
+
# bearer for the rest of a session is authority nothing needs.
|
|
1763
|
+
client.get_httpx_client().close()
|
|
1764
|
+
lease, credential = parse_session_lease(wire)
|
|
1765
|
+
if lease.session_id != invocation.session_id or str(lease.workspace_id) != (
|
|
1766
|
+
invocation.workspace_id
|
|
1767
|
+
):
|
|
1768
|
+
credential.close()
|
|
1769
|
+
raise SessionWorkerError(
|
|
1770
|
+
"SESSION_LEASE_MISMATCH", "the lease names a different session or workspace"
|
|
1771
|
+
)
|
|
1772
|
+
return lease, credential
|
|
1773
|
+
|
|
1774
|
+
|
|
1775
|
+
# ==========================================================================================
|
|
1776
|
+
# The supervisor
|
|
1777
|
+
# ==========================================================================================
|
|
1778
|
+
|
|
1779
|
+
|
|
1780
|
+
@dataclass
|
|
1781
|
+
class _SessionOutcome:
|
|
1782
|
+
session_id: str
|
|
1783
|
+
attempt_id: str | None = None
|
|
1784
|
+
probes_answered: int = 0
|
|
1785
|
+
probes_failed: int = 0
|
|
1786
|
+
terminal_state: str | None = None
|
|
1787
|
+
failure_code: str | None = None
|
|
1788
|
+
warm_pool_ready: int = 0
|
|
1789
|
+
catalogue: Mapping[str, Any] | None = None
|
|
1790
|
+
events: list[str] = field(default_factory=list)
|
|
1791
|
+
|
|
1792
|
+
def to_dict(self) -> dict[str, Any]:
|
|
1793
|
+
return {
|
|
1794
|
+
"status": "failed" if self.failure_code else "stopped",
|
|
1795
|
+
"role": SESSION_WORKER_ROLE,
|
|
1796
|
+
"session_id": self.session_id,
|
|
1797
|
+
"attempt_id": self.attempt_id,
|
|
1798
|
+
"probes_answered": self.probes_answered,
|
|
1799
|
+
"probes_failed": self.probes_failed,
|
|
1800
|
+
"terminal_state": self.terminal_state,
|
|
1801
|
+
"failure_code": self.failure_code,
|
|
1802
|
+
"warm_pool_ready": self.warm_pool_ready,
|
|
1803
|
+
}
|
|
1804
|
+
|
|
1805
|
+
|
|
1806
|
+
class _Heartbeat(threading.Thread):
|
|
1807
|
+
"""Liveness on a timer, and nothing else.
|
|
1808
|
+
|
|
1809
|
+
A separate thread rather than a tick inside the claim loop, because the claim loop blocks for
|
|
1810
|
+
the length of a probe and a probe may run for the better part of a minute. A heartbeat that
|
|
1811
|
+
only fires between probes would go quiet during exactly the interval where a client most wants
|
|
1812
|
+
to know the worker is alive.
|
|
1813
|
+
|
|
1814
|
+
Every failure is swallowed. A missed heartbeat costs liveness resolution; the authority on
|
|
1815
|
+
whether this worker still holds its lease is the attempt lease expiry Studio enforces on every
|
|
1816
|
+
fenced call, and the claim loop discovers a lost lease there. Failing the session from this
|
|
1817
|
+
thread would be this thread deciding something it cannot see.
|
|
1818
|
+
"""
|
|
1819
|
+
|
|
1820
|
+
def __init__(
|
|
1821
|
+
self,
|
|
1822
|
+
*,
|
|
1823
|
+
send: Callable[[], None],
|
|
1824
|
+
interval_seconds: float,
|
|
1825
|
+
stop: threading.Event,
|
|
1826
|
+
deadline: SessionDeadline | None = None,
|
|
1827
|
+
on_expired: Callable[[], None] | None = None,
|
|
1828
|
+
) -> None:
|
|
1829
|
+
super().__init__(name="mr-session-heartbeat", daemon=True)
|
|
1830
|
+
self._send = send
|
|
1831
|
+
self._interval = max(1.0, float(interval_seconds))
|
|
1832
|
+
self._stop = stop
|
|
1833
|
+
self._deadline = deadline
|
|
1834
|
+
self._on_expired = on_expired
|
|
1835
|
+
self.sent = 0
|
|
1836
|
+
|
|
1837
|
+
def run(self) -> None:
|
|
1838
|
+
while not self._stop.is_set():
|
|
1839
|
+
wait_seconds = self._interval
|
|
1840
|
+
if self._deadline is not None:
|
|
1841
|
+
wait_seconds = self._deadline.bounded_wait(wait_seconds)
|
|
1842
|
+
if wait_seconds <= 0:
|
|
1843
|
+
self._expire()
|
|
1844
|
+
return
|
|
1845
|
+
if self._stop.wait(wait_seconds):
|
|
1846
|
+
return
|
|
1847
|
+
if self._deadline is not None and self._deadline.expired():
|
|
1848
|
+
self._expire()
|
|
1849
|
+
return
|
|
1850
|
+
try:
|
|
1851
|
+
self._send()
|
|
1852
|
+
except Exception:
|
|
1853
|
+
LOGGER.debug("session heartbeat failed", exc_info=True)
|
|
1854
|
+
continue
|
|
1855
|
+
self.sent += 1
|
|
1856
|
+
|
|
1857
|
+
def _expire(self) -> None:
|
|
1858
|
+
self._stop.set()
|
|
1859
|
+
if self._on_expired is not None:
|
|
1860
|
+
self._on_expired()
|
|
1861
|
+
|
|
1862
|
+
|
|
1863
|
+
class SessionSupervisor:
|
|
1864
|
+
"""Lease, start, claim, answer, heartbeat, and stop on a terminal state.
|
|
1865
|
+
|
|
1866
|
+
In production Job mode, a capability-free parent supplies an opaque dispatch and this
|
|
1867
|
+
supervisor exchanges it only in its isolated child. The capability remains in that child
|
|
1868
|
+
inside ``CredentialLease`` only for the life of the session; it is not serialised, logged,
|
|
1869
|
+
written to disk or passed to the probe sandbox, and it is zeroed when the supervisor stops.
|
|
1870
|
+
``invocation``/``nonce`` remain only as a compatibility seam for the retained rollback
|
|
1871
|
+
Service; Job mode never constructs or accepts a lease nonce.
|
|
1872
|
+
"""
|
|
1873
|
+
|
|
1874
|
+
def __init__(
|
|
1875
|
+
self,
|
|
1876
|
+
*,
|
|
1877
|
+
config: SessionWorkerConfig,
|
|
1878
|
+
invocation: StartInvocation | None = None,
|
|
1879
|
+
nonce: bytearray | None = None,
|
|
1880
|
+
bootstrap: tuple[SessionLease, CredentialLease] | None = None,
|
|
1881
|
+
dispatch: SessionDispatch | None = None,
|
|
1882
|
+
transport: BootstrapTransport | None = None,
|
|
1883
|
+
sandbox_factory: Callable[[SessionWorkerConfig], Any] | None = None,
|
|
1884
|
+
reader_factory: Callable[[SessionLease], ProbeSourceReader] | None = None,
|
|
1885
|
+
lease_session: Callable[..., tuple[SessionLease, CredentialLease]] = _lease_session,
|
|
1886
|
+
sleep: Callable[[float], None] = time.sleep,
|
|
1887
|
+
wall_clock: Callable[[], datetime] = _wall_clock_now,
|
|
1888
|
+
monotonic_clock: Callable[[], float] = time.monotonic,
|
|
1889
|
+
process_isolation: bool = False,
|
|
1890
|
+
on_isolation_ready: Callable[[], None] | None = None,
|
|
1891
|
+
) -> None:
|
|
1892
|
+
legacy = invocation is not None or nonce is not None
|
|
1893
|
+
modes = int(legacy) + int(bootstrap is not None) + int(dispatch is not None)
|
|
1894
|
+
if modes != 1 or (legacy and (invocation is None or nonce is None)):
|
|
1895
|
+
raise SessionWorkerError(
|
|
1896
|
+
"SESSION_SUPERVISOR_CONFIGURATION_INVALID",
|
|
1897
|
+
"provide exactly one legacy invocation, dispatch, or dispatch bootstrap",
|
|
1898
|
+
)
|
|
1899
|
+
self._config = config
|
|
1900
|
+
self._invocation = invocation
|
|
1901
|
+
self._nonce = nonce
|
|
1902
|
+
self._bootstrap = bootstrap
|
|
1903
|
+
self._dispatch = dispatch
|
|
1904
|
+
self._transport = transport or UrlLibBootstrapTransport()
|
|
1905
|
+
self._sandbox_factory = sandbox_factory or _default_sandbox
|
|
1906
|
+
# `None` rather than a default lambda, because the production reader cannot be built from
|
|
1907
|
+
# a lease alone: it needs the fenced authority and a transfer client, both of which exist
|
|
1908
|
+
# only once `_serve` has one. A test injecting a factory still overrides everything.
|
|
1909
|
+
self._reader_factory = reader_factory
|
|
1910
|
+
self._lease_session = lease_session
|
|
1911
|
+
self._sleep = sleep
|
|
1912
|
+
self._wall_clock = wall_clock
|
|
1913
|
+
self._monotonic_clock = monotonic_clock
|
|
1914
|
+
self._process_isolation = process_isolation
|
|
1915
|
+
self._on_isolation_ready = on_isolation_ready
|
|
1916
|
+
self._stop = threading.Event()
|
|
1917
|
+
self.outcome = _SessionOutcome(
|
|
1918
|
+
session_id=(
|
|
1919
|
+
bootstrap[0].session_id
|
|
1920
|
+
if bootstrap is not None
|
|
1921
|
+
else invocation.session_id
|
|
1922
|
+
if invocation is not None
|
|
1923
|
+
else ""
|
|
1924
|
+
)
|
|
1925
|
+
)
|
|
1926
|
+
|
|
1927
|
+
def stop(self) -> None:
|
|
1928
|
+
"""Ask the claim loop to finish the probe it is on and then stop."""
|
|
1929
|
+
|
|
1930
|
+
self._stop.set()
|
|
1931
|
+
|
|
1932
|
+
@property
|
|
1933
|
+
def is_running(self) -> bool:
|
|
1934
|
+
"""Whether this supervisor still owns a live session."""
|
|
1935
|
+
|
|
1936
|
+
return not self._stop.is_set() and self.outcome.terminal_state is None
|
|
1937
|
+
|
|
1938
|
+
def _stop_at_deadline(self, deadline: SessionDeadline) -> bool:
|
|
1939
|
+
"""Settle the local process at the credential boundary, without a producer callback."""
|
|
1940
|
+
|
|
1941
|
+
if not deadline.expired():
|
|
1942
|
+
return False
|
|
1943
|
+
if self.outcome.terminal_state is None:
|
|
1944
|
+
self.outcome.terminal_state = "capability_expired"
|
|
1945
|
+
self._stop.set()
|
|
1946
|
+
return True
|
|
1947
|
+
|
|
1948
|
+
# -- the run ----------------------------------------------------------------------------
|
|
1949
|
+
|
|
1950
|
+
def run(self) -> dict[str, Any]:
|
|
1951
|
+
"""Own one session end to end. Never raises; the outcome carries the failure code."""
|
|
1952
|
+
|
|
1953
|
+
if self._process_isolation:
|
|
1954
|
+
return self._run_isolated()
|
|
1955
|
+
return self._run_in_process()
|
|
1956
|
+
|
|
1957
|
+
def _run_isolated(self) -> dict[str, Any]:
|
|
1958
|
+
"""Run the capability-bearing half in a child and enforce its expiry from this parent.
|
|
1959
|
+
|
|
1960
|
+
The parent receives only a session ID and absolute expiry, never the capability. It is
|
|
1961
|
+
therefore able to wait, kill and return exit-0 semantics even if a reader or native parser
|
|
1962
|
+
in the child never returns to Python or holds the child GIL. The child arms its own
|
|
1963
|
+
default-action ``SIGALRM`` before it acknowledges the lease, closing the small scheduling
|
|
1964
|
+
gap before this parent has installed its paired-clock monitor.
|
|
1965
|
+
"""
|
|
1966
|
+
|
|
1967
|
+
if self._bootstrap is not None:
|
|
1968
|
+
raise SessionWorkerError(
|
|
1969
|
+
"SESSION_SUPERVISOR_CONFIGURATION_INVALID",
|
|
1970
|
+
"a capability-bearing bootstrap cannot enter process isolation",
|
|
1971
|
+
)
|
|
1972
|
+
if not hasattr(os, "fork") or not hasattr(signal, "setitimer"):
|
|
1973
|
+
raise SessionWorkerError(
|
|
1974
|
+
"SESSION_PROCESS_ISOLATION_UNAVAILABLE",
|
|
1975
|
+
"the session deadline requires POSIX process and interval-timer support",
|
|
1976
|
+
)
|
|
1977
|
+
previous_sigterm: Any | None = None
|
|
1978
|
+
if threading.current_thread() is threading.main_thread():
|
|
1979
|
+
previous_sigterm = signal.getsignal(signal.SIGTERM)
|
|
1980
|
+
|
|
1981
|
+
def request_stop(_signum: int, _frame: Any) -> None:
|
|
1982
|
+
# The monitor owns no capability. Do only the async-safe state change here; the
|
|
1983
|
+
# ordinary monitor path kills/reaps every tracked process group.
|
|
1984
|
+
self._stop.set()
|
|
1985
|
+
|
|
1986
|
+
# Install before the fork. A rejected/timeout legacy bootstrap must be able to signal
|
|
1987
|
+
# its outer monitor even if its inner child has already made a separate process group.
|
|
1988
|
+
signal.signal(signal.SIGTERM, request_stop)
|
|
1989
|
+
prior_subreaper = False
|
|
1990
|
+
subreaper_enabled = False
|
|
1991
|
+
baseline_children = frozenset()
|
|
1992
|
+
if sys.platform == "linux":
|
|
1993
|
+
try:
|
|
1994
|
+
prior_subreaper = _child_subreaper_enabled()
|
|
1995
|
+
_set_child_subreaper(True)
|
|
1996
|
+
subreaper_enabled = True
|
|
1997
|
+
# The monitor is dedicated to one session. A baseline still makes its teardown
|
|
1998
|
+
# exact if an embedding process happened to own a child before supervision began.
|
|
1999
|
+
baseline_children = _linux_direct_child_pids()
|
|
2000
|
+
except OSError as error:
|
|
2001
|
+
if subreaper_enabled:
|
|
2002
|
+
with contextlib.suppress(OSError):
|
|
2003
|
+
_set_child_subreaper(prior_subreaper)
|
|
2004
|
+
if previous_sigterm is not None:
|
|
2005
|
+
signal.signal(signal.SIGTERM, previous_sigterm)
|
|
2006
|
+
raise SessionWorkerError(
|
|
2007
|
+
"SESSION_PROCESS_ISOLATION_UNAVAILABLE",
|
|
2008
|
+
"could not establish isolated-session descendant supervision",
|
|
2009
|
+
) from error
|
|
2010
|
+
try:
|
|
2011
|
+
child_to_parent_read, child_to_parent_write = os.pipe()
|
|
2012
|
+
parent_to_child_read, parent_to_child_write = os.pipe()
|
|
2013
|
+
except OSError as error:
|
|
2014
|
+
if subreaper_enabled:
|
|
2015
|
+
with contextlib.suppress(OSError):
|
|
2016
|
+
_set_child_subreaper(prior_subreaper)
|
|
2017
|
+
if previous_sigterm is not None:
|
|
2018
|
+
signal.signal(signal.SIGTERM, previous_sigterm)
|
|
2019
|
+
raise SessionWorkerError(
|
|
2020
|
+
"SESSION_PROCESS_ISOLATION_UNAVAILABLE",
|
|
2021
|
+
"could not allocate isolated-session control pipes",
|
|
2022
|
+
) from error
|
|
2023
|
+
if self._stop.is_set():
|
|
2024
|
+
for descriptor in (
|
|
2025
|
+
child_to_parent_read,
|
|
2026
|
+
child_to_parent_write,
|
|
2027
|
+
parent_to_child_read,
|
|
2028
|
+
parent_to_child_write,
|
|
2029
|
+
):
|
|
2030
|
+
with contextlib.suppress(OSError):
|
|
2031
|
+
os.close(descriptor)
|
|
2032
|
+
if subreaper_enabled:
|
|
2033
|
+
with contextlib.suppress(OSError):
|
|
2034
|
+
_set_child_subreaper(prior_subreaper)
|
|
2035
|
+
if previous_sigterm is not None:
|
|
2036
|
+
signal.signal(signal.SIGTERM, previous_sigterm)
|
|
2037
|
+
self.outcome.terminal_state = "service_stopped"
|
|
2038
|
+
return self.outcome.to_dict()
|
|
2039
|
+
try:
|
|
2040
|
+
child_pid = os.fork()
|
|
2041
|
+
except OSError as error:
|
|
2042
|
+
for descriptor in (
|
|
2043
|
+
child_to_parent_read,
|
|
2044
|
+
child_to_parent_write,
|
|
2045
|
+
parent_to_child_read,
|
|
2046
|
+
parent_to_child_write,
|
|
2047
|
+
):
|
|
2048
|
+
os.close(descriptor)
|
|
2049
|
+
if subreaper_enabled:
|
|
2050
|
+
with contextlib.suppress(OSError):
|
|
2051
|
+
_set_child_subreaper(prior_subreaper)
|
|
2052
|
+
if previous_sigterm is not None:
|
|
2053
|
+
signal.signal(signal.SIGTERM, previous_sigterm)
|
|
2054
|
+
raise SessionWorkerError(
|
|
2055
|
+
"SESSION_PROCESS_ISOLATION_UNAVAILABLE", "could not start an isolated session child"
|
|
2056
|
+
) from error
|
|
2057
|
+
if child_pid == 0: # pragma: no cover - exercised through the parent process contract.
|
|
2058
|
+
if previous_sigterm is not None:
|
|
2059
|
+
signal.signal(signal.SIGTERM, previous_sigterm)
|
|
2060
|
+
os.close(child_to_parent_read)
|
|
2061
|
+
os.close(parent_to_child_write)
|
|
2062
|
+
self._process_isolation = False
|
|
2063
|
+
try:
|
|
2064
|
+
# Give this one capability-bearing process its own group before it can create a
|
|
2065
|
+
# sandbox. The parent kills this group at expiry; separate sandbox groups are
|
|
2066
|
+
# reparented to the Linux subreaper and killed/reaped there.
|
|
2067
|
+
os.setsid()
|
|
2068
|
+
summary = self._run_in_process(
|
|
2069
|
+
on_lease=lambda lease, deadline: self._announce_isolated_lease(
|
|
2070
|
+
child_to_parent_write, parent_to_child_read, lease, deadline
|
|
2071
|
+
)
|
|
2072
|
+
)
|
|
2073
|
+
self._write_isolated_message(
|
|
2074
|
+
child_to_parent_write, {"kind": "summary", "summary": summary}
|
|
2075
|
+
)
|
|
2076
|
+
os._exit(0)
|
|
2077
|
+
except BaseException as error:
|
|
2078
|
+
# The parent gets a bounded, non-secret diagnostic instead of a traceback that
|
|
2079
|
+
# could retain the capability-bearing stack. ``_run_in_process`` normally
|
|
2080
|
+
# turns operational failures into a summary; this is only its own control plane.
|
|
2081
|
+
with contextlib.suppress(Exception):
|
|
2082
|
+
self._write_isolated_message(
|
|
2083
|
+
child_to_parent_write,
|
|
2084
|
+
{
|
|
2085
|
+
"kind": "error",
|
|
2086
|
+
"code": getattr(error, "code", "SESSION_CHILD_FAILED"),
|
|
2087
|
+
},
|
|
2088
|
+
)
|
|
2089
|
+
os._exit(1)
|
|
2090
|
+
finally:
|
|
2091
|
+
with contextlib.suppress(OSError):
|
|
2092
|
+
os.close(child_to_parent_write)
|
|
2093
|
+
with contextlib.suppress(OSError):
|
|
2094
|
+
os.close(parent_to_child_read)
|
|
2095
|
+
|
|
2096
|
+
os.close(child_to_parent_write)
|
|
2097
|
+
os.close(parent_to_child_read)
|
|
2098
|
+
# A fork duplicates a legacy nonce but not a capability (the child performs exchange).
|
|
2099
|
+
# Retire the parent's copy before it waits so the service process cannot retain it.
|
|
2100
|
+
if self._nonce is not None:
|
|
2101
|
+
self._nonce[:] = b"\x00" * len(self._nonce)
|
|
2102
|
+
try:
|
|
2103
|
+
if self._stop.is_set():
|
|
2104
|
+
self.outcome.terminal_state = "service_stopped"
|
|
2105
|
+
return self._terminate_isolated_child(child_pid)
|
|
2106
|
+
if self._on_isolation_ready is not None:
|
|
2107
|
+
try:
|
|
2108
|
+
self._on_isolation_ready()
|
|
2109
|
+
except Exception as error:
|
|
2110
|
+
self.outcome.failure_code = "SESSION_ISOLATION_READY_FAILED"
|
|
2111
|
+
LOGGER.warning(
|
|
2112
|
+
"isolated-session readiness acknowledgement failed: %s",
|
|
2113
|
+
type(error).__name__,
|
|
2114
|
+
)
|
|
2115
|
+
return self._terminate_isolated_child(child_pid)
|
|
2116
|
+
return self._monitor_isolated_child(
|
|
2117
|
+
child_pid, child_to_parent_read, parent_to_child_write
|
|
2118
|
+
)
|
|
2119
|
+
finally:
|
|
2120
|
+
if previous_sigterm is not None:
|
|
2121
|
+
signal.signal(signal.SIGTERM, previous_sigterm)
|
|
2122
|
+
with contextlib.suppress(OSError):
|
|
2123
|
+
os.close(child_to_parent_read)
|
|
2124
|
+
with contextlib.suppress(OSError):
|
|
2125
|
+
os.close(parent_to_child_write)
|
|
2126
|
+
if subreaper_enabled:
|
|
2127
|
+
try:
|
|
2128
|
+
self._kill_adopted_session_descendants(baseline_children)
|
|
2129
|
+
finally:
|
|
2130
|
+
_set_child_subreaper(prior_subreaper)
|
|
2131
|
+
|
|
2132
|
+
@staticmethod
|
|
2133
|
+
def _write_isolated_message(descriptor: int, payload: dict[str, Any]) -> None:
|
|
2134
|
+
raw = canonical_json_bytes(payload) + b"\n"
|
|
2135
|
+
if len(raw) > MAX_SESSION_SUPERVISOR_MESSAGE_BYTES:
|
|
2136
|
+
raise SessionWorkerError(
|
|
2137
|
+
"SESSION_CHILD_PROTOCOL_INVALID", "isolated-session message exceeds its byte limit"
|
|
2138
|
+
)
|
|
2139
|
+
view = memoryview(raw)
|
|
2140
|
+
try:
|
|
2141
|
+
while view:
|
|
2142
|
+
written = os.write(descriptor, view)
|
|
2143
|
+
if written <= 0:
|
|
2144
|
+
raise OSError("short isolated-session pipe write")
|
|
2145
|
+
view = view[written:]
|
|
2146
|
+
finally:
|
|
2147
|
+
view.release()
|
|
2148
|
+
|
|
2149
|
+
@staticmethod
|
|
2150
|
+
def _read_isolated_message(
|
|
2151
|
+
descriptor: int, buffer: bytearray, *, timeout: float | None
|
|
2152
|
+
) -> dict[str, Any] | None:
|
|
2153
|
+
readable, _unused, _errors = select.select([descriptor], [], [], timeout)
|
|
2154
|
+
if not readable:
|
|
2155
|
+
return None
|
|
2156
|
+
chunk = os.read(descriptor, 64 * 1024)
|
|
2157
|
+
if not chunk:
|
|
2158
|
+
raise EOFError("isolated session child closed its control pipe")
|
|
2159
|
+
buffer.extend(chunk)
|
|
2160
|
+
if len(buffer) > MAX_SESSION_SUPERVISOR_MESSAGE_BYTES:
|
|
2161
|
+
raise SessionWorkerError(
|
|
2162
|
+
"SESSION_CHILD_PROTOCOL_INVALID", "isolated-session message exceeds its byte limit"
|
|
2163
|
+
)
|
|
2164
|
+
newline = buffer.find(b"\n")
|
|
2165
|
+
if newline < 0:
|
|
2166
|
+
return None
|
|
2167
|
+
raw = bytes(buffer[:newline])
|
|
2168
|
+
del buffer[: newline + 1]
|
|
2169
|
+
value = _strict_json(raw, "isolated-session message")
|
|
2170
|
+
if not isinstance(value, dict):
|
|
2171
|
+
raise SessionWorkerError(
|
|
2172
|
+
"SESSION_CHILD_PROTOCOL_INVALID", "isolated-session message is not an object"
|
|
2173
|
+
)
|
|
2174
|
+
return value
|
|
2175
|
+
|
|
2176
|
+
def _announce_isolated_lease(
|
|
2177
|
+
self,
|
|
2178
|
+
write_descriptor: int,
|
|
2179
|
+
read_descriptor: int,
|
|
2180
|
+
lease: SessionLease,
|
|
2181
|
+
deadline: SessionDeadline,
|
|
2182
|
+
) -> None:
|
|
2183
|
+
remaining = deadline.remaining_seconds()
|
|
2184
|
+
if remaining <= 0:
|
|
2185
|
+
raise SessionWorkerError(
|
|
2186
|
+
"SESSION_CAPABILITY_EXPIRED", "the session capability expired before supervision"
|
|
2187
|
+
)
|
|
2188
|
+
# SIG_DFL is a kernel-enforced deadline: a child blocked in native code cannot delay it.
|
|
2189
|
+
signal.signal(signal.SIGALRM, signal.SIG_DFL)
|
|
2190
|
+
signal.setitimer(signal.ITIMER_REAL, remaining)
|
|
2191
|
+
expiry = lease.capability_expires_at.isoformat(timespec="microseconds").replace(
|
|
2192
|
+
"+00:00", "Z"
|
|
2193
|
+
)
|
|
2194
|
+
self._write_isolated_message(
|
|
2195
|
+
write_descriptor,
|
|
2196
|
+
{"kind": "lease", "session_id": lease.session_id, "capability_expires_at": expiry},
|
|
2197
|
+
)
|
|
2198
|
+
response = bytearray()
|
|
2199
|
+
while True:
|
|
2200
|
+
command = self._read_isolated_message(read_descriptor, response, timeout=None)
|
|
2201
|
+
if command is not None:
|
|
2202
|
+
break
|
|
2203
|
+
if command != {"kind": "start"}:
|
|
2204
|
+
raise SessionWorkerError(
|
|
2205
|
+
"SESSION_CHILD_PROTOCOL_INVALID", "isolated-session parent did not authorize start"
|
|
2206
|
+
)
|
|
2207
|
+
|
|
2208
|
+
def _monitor_isolated_child(
|
|
2209
|
+
self, child_pid: int, read_descriptor: int, write_descriptor: int
|
|
2210
|
+
) -> dict[str, Any]:
|
|
2211
|
+
buffer = bytearray()
|
|
2212
|
+
deadline: SessionDeadline | None = None
|
|
2213
|
+
while True:
|
|
2214
|
+
if self._stop.is_set():
|
|
2215
|
+
if self.outcome.terminal_state is None:
|
|
2216
|
+
self.outcome.terminal_state = "service_stopped"
|
|
2217
|
+
return self._terminate_isolated_child(child_pid)
|
|
2218
|
+
timeout = 1.0 if deadline is None else deadline.bounded_wait(1.0)
|
|
2219
|
+
if deadline is not None and timeout <= 0:
|
|
2220
|
+
return self._expire_isolated_child(child_pid)
|
|
2221
|
+
try:
|
|
2222
|
+
message = self._read_isolated_message(read_descriptor, buffer, timeout=timeout)
|
|
2223
|
+
except EOFError:
|
|
2224
|
+
status = self._wait_for_isolated_child(child_pid)
|
|
2225
|
+
if deadline is not None and (
|
|
2226
|
+
deadline.expired()
|
|
2227
|
+
or (os.WIFSIGNALED(status) and os.WTERMSIG(status) == signal.SIGALRM)
|
|
2228
|
+
):
|
|
2229
|
+
return self._deadline_summary()
|
|
2230
|
+
self.outcome.failure_code = "SESSION_CHILD_EXITED"
|
|
2231
|
+
return self.outcome.to_dict()
|
|
2232
|
+
if message is None:
|
|
2233
|
+
continue
|
|
2234
|
+
kind = message.get("kind")
|
|
2235
|
+
if kind == "lease":
|
|
2236
|
+
if deadline is not None or set(message) != {
|
|
2237
|
+
"kind",
|
|
2238
|
+
"session_id",
|
|
2239
|
+
"capability_expires_at",
|
|
2240
|
+
}:
|
|
2241
|
+
self.outcome.failure_code = "SESSION_CHILD_PROTOCOL_INVALID"
|
|
2242
|
+
return self._terminate_isolated_child(child_pid)
|
|
2243
|
+
try:
|
|
2244
|
+
session_id = _uuid_text(message["session_id"], "isolated session id")
|
|
2245
|
+
expires_at = _parse_utc_timestamp(
|
|
2246
|
+
message["capability_expires_at"], "isolated capability_expires_at"
|
|
2247
|
+
)
|
|
2248
|
+
deadline = SessionDeadline.from_expiry(
|
|
2249
|
+
expires_at,
|
|
2250
|
+
wall_clock=self._wall_clock,
|
|
2251
|
+
monotonic_clock=self._monotonic_clock,
|
|
2252
|
+
)
|
|
2253
|
+
except SessionWorkerError as error:
|
|
2254
|
+
if error.code == "SESSION_CAPABILITY_EXPIRED":
|
|
2255
|
+
return self._expire_isolated_child(child_pid)
|
|
2256
|
+
self.outcome.failure_code = error.code
|
|
2257
|
+
return self._terminate_isolated_child(child_pid)
|
|
2258
|
+
self.outcome.session_id = session_id
|
|
2259
|
+
self._write_isolated_message(write_descriptor, {"kind": "start"})
|
|
2260
|
+
continue
|
|
2261
|
+
if kind == "summary" and set(message) == {"kind", "summary"}:
|
|
2262
|
+
status = self._wait_for_isolated_child(child_pid)
|
|
2263
|
+
summary = message["summary"]
|
|
2264
|
+
if not isinstance(summary, dict):
|
|
2265
|
+
self.outcome.failure_code = "SESSION_CHILD_PROTOCOL_INVALID"
|
|
2266
|
+
return self.outcome.to_dict()
|
|
2267
|
+
if deadline is not None and (
|
|
2268
|
+
deadline.expired()
|
|
2269
|
+
or (os.WIFSIGNALED(status) and os.WTERMSIG(status) == signal.SIGALRM)
|
|
2270
|
+
):
|
|
2271
|
+
return self._deadline_summary()
|
|
2272
|
+
return summary
|
|
2273
|
+
if (
|
|
2274
|
+
kind == "error"
|
|
2275
|
+
and set(message) == {"kind", "code"}
|
|
2276
|
+
and isinstance(message["code"], str)
|
|
2277
|
+
):
|
|
2278
|
+
self.outcome.failure_code = message["code"]
|
|
2279
|
+
else:
|
|
2280
|
+
self.outcome.failure_code = "SESSION_CHILD_PROTOCOL_INVALID"
|
|
2281
|
+
return self._terminate_isolated_child(child_pid)
|
|
2282
|
+
|
|
2283
|
+
def _expire_isolated_child(self, child_pid: int) -> dict[str, Any]:
|
|
2284
|
+
self._kill_isolated_child_group(child_pid)
|
|
2285
|
+
self._wait_for_isolated_child(child_pid)
|
|
2286
|
+
return self._deadline_summary()
|
|
2287
|
+
|
|
2288
|
+
def _terminate_isolated_child(self, child_pid: int) -> dict[str, Any]:
|
|
2289
|
+
self._kill_isolated_child_group(child_pid)
|
|
2290
|
+
self._wait_for_isolated_child(child_pid)
|
|
2291
|
+
return self.outcome.to_dict()
|
|
2292
|
+
|
|
2293
|
+
@staticmethod
|
|
2294
|
+
def _kill_isolated_child_group(child_pid: int) -> None:
|
|
2295
|
+
"""Kill the child process group, with a PID fallback before ``setsid`` can complete."""
|
|
2296
|
+
|
|
2297
|
+
try:
|
|
2298
|
+
os.killpg(child_pid, signal.SIGKILL)
|
|
2299
|
+
return
|
|
2300
|
+
except (PermissionError, ProcessLookupError):
|
|
2301
|
+
# Before the child's first control message it may not yet have made itself a group
|
|
2302
|
+
# leader. Some Darwin process-group races also report EPERM after the group leader has
|
|
2303
|
+
# exited. Killing the exact direct PID still prevents a bootstrap/protocol error from
|
|
2304
|
+
# leaving the capability exchange running.
|
|
2305
|
+
pass
|
|
2306
|
+
with contextlib.suppress(PermissionError, ProcessLookupError):
|
|
2307
|
+
os.kill(child_pid, signal.SIGKILL)
|
|
2308
|
+
|
|
2309
|
+
@staticmethod
|
|
2310
|
+
def _wait_for_isolated_child(child_pid: int) -> int:
|
|
2311
|
+
while True:
|
|
2312
|
+
try:
|
|
2313
|
+
_pid, status = os.waitpid(child_pid, 0)
|
|
2314
|
+
return status
|
|
2315
|
+
except InterruptedError:
|
|
2316
|
+
continue
|
|
2317
|
+
except ChildProcessError:
|
|
2318
|
+
return 0
|
|
2319
|
+
|
|
2320
|
+
@staticmethod
|
|
2321
|
+
def _kill_adopted_session_descendants(baseline_children: frozenset[int]) -> None:
|
|
2322
|
+
"""Kill/reap Linux descendants adopted after the capability child dies.
|
|
2323
|
+
|
|
2324
|
+
Sandboxes intentionally create separate process groups so parser helpers cannot affect the
|
|
2325
|
+
coordinator. A force-killed capability child cannot run ``close_warm_pool``; as a Linux
|
|
2326
|
+
subreaper, this parent adopts those groups and kills their live leader/group before each
|
|
2327
|
+
reap. The monitor creates no other children, and the pre-fork baseline prevents touching
|
|
2328
|
+
a host process's pre-existing children if this code is embedded in a test or adapter.
|
|
2329
|
+
"""
|
|
2330
|
+
|
|
2331
|
+
if sys.platform != "linux":
|
|
2332
|
+
return
|
|
2333
|
+
deadline = time.monotonic() + 5.0
|
|
2334
|
+
own_group = os.getpgrp()
|
|
2335
|
+
while True:
|
|
2336
|
+
children = _linux_direct_child_pids() - baseline_children
|
|
2337
|
+
if not children:
|
|
2338
|
+
return
|
|
2339
|
+
for child_pid in children:
|
|
2340
|
+
try:
|
|
2341
|
+
group = os.getpgid(child_pid)
|
|
2342
|
+
except ProcessLookupError:
|
|
2343
|
+
continue
|
|
2344
|
+
try:
|
|
2345
|
+
if group == own_group:
|
|
2346
|
+
# The dedicated capability child makes its own group before any sandbox
|
|
2347
|
+
# launch. A shared group is therefore never safe to signal wholesale.
|
|
2348
|
+
os.kill(child_pid, signal.SIGKILL)
|
|
2349
|
+
else:
|
|
2350
|
+
os.killpg(group, signal.SIGKILL)
|
|
2351
|
+
except ProcessLookupError:
|
|
2352
|
+
continue
|
|
2353
|
+
for child_pid in children:
|
|
2354
|
+
with contextlib.suppress(ChildProcessError):
|
|
2355
|
+
os.waitpid(child_pid, os.WNOHANG)
|
|
2356
|
+
if time.monotonic() >= deadline:
|
|
2357
|
+
raise SessionWorkerError(
|
|
2358
|
+
"SESSION_DESCENDANT_CLEANUP_FAILED",
|
|
2359
|
+
"isolated-session descendants did not exit after the capability child",
|
|
2360
|
+
)
|
|
2361
|
+
time.sleep(0.005)
|
|
2362
|
+
|
|
2363
|
+
def _deadline_summary(self) -> dict[str, Any]:
|
|
2364
|
+
self.outcome.terminal_state = "capability_expired"
|
|
2365
|
+
self._stop.set()
|
|
2366
|
+
return self.outcome.to_dict()
|
|
2367
|
+
|
|
2368
|
+
def _run_in_process(
|
|
2369
|
+
self, *, on_lease: Callable[[SessionLease, SessionDeadline], None] | None = None
|
|
2370
|
+
) -> dict[str, Any]:
|
|
2371
|
+
"""Run the capability-bearing half, optionally after an isolated-parent handshake."""
|
|
2372
|
+
|
|
2373
|
+
if self._bootstrap is not None:
|
|
2374
|
+
# Transfer ownership exactly once. A duplicate Job execution must obtain a fresh
|
|
2375
|
+
# bootstrap through Studio; it cannot re-run this supervisor with an old capability.
|
|
2376
|
+
lease, credential = self._bootstrap
|
|
2377
|
+
self._bootstrap = None
|
|
2378
|
+
elif self._dispatch is not None:
|
|
2379
|
+
try:
|
|
2380
|
+
lease, credential = exchange_session_dispatch(
|
|
2381
|
+
config=self._config, dispatch=self._dispatch, transport=self._transport
|
|
2382
|
+
)
|
|
2383
|
+
except (SessionWorkerError, HostedWorkerError) as error:
|
|
2384
|
+
code = getattr(error, "code", "SESSION_LEASE_FAILED")
|
|
2385
|
+
if code in CONSUMED_SESSION_DISPATCH_CODES:
|
|
2386
|
+
self.outcome.terminal_state = "dispatch_consumed"
|
|
2387
|
+
else:
|
|
2388
|
+
self.outcome.failure_code = code
|
|
2389
|
+
return self.outcome.to_dict()
|
|
2390
|
+
else:
|
|
2391
|
+
assert self._invocation is not None
|
|
2392
|
+
assert self._nonce is not None
|
|
2393
|
+
try:
|
|
2394
|
+
lease, credential = self._lease_session(
|
|
2395
|
+
config=self._config,
|
|
2396
|
+
invocation=self._invocation,
|
|
2397
|
+
nonce=self._nonce,
|
|
2398
|
+
transport=self._transport,
|
|
2399
|
+
)
|
|
2400
|
+
except (SessionWorkerError, HostedWorkerError) as error:
|
|
2401
|
+
self.outcome.failure_code = getattr(error, "code", "SESSION_LEASE_FAILED")
|
|
2402
|
+
return self.outcome.to_dict()
|
|
2403
|
+
finally:
|
|
2404
|
+
# Zeroed the instant the legacy lease exchange is over, success or failure. One
|
|
2405
|
+
# use, one lifetime. Job mode has no nonce at all.
|
|
2406
|
+
self._nonce[:] = b"\x00" * len(self._nonce)
|
|
2407
|
+
self.outcome.attempt_id = str(lease.attempt_id)
|
|
2408
|
+
try:
|
|
2409
|
+
with credential:
|
|
2410
|
+
deadline = SessionDeadline.from_lease(
|
|
2411
|
+
lease,
|
|
2412
|
+
wall_clock=self._wall_clock,
|
|
2413
|
+
monotonic_clock=self._monotonic_clock,
|
|
2414
|
+
)
|
|
2415
|
+
if on_lease is not None:
|
|
2416
|
+
on_lease(lease, deadline)
|
|
2417
|
+
try:
|
|
2418
|
+
self.outcome.catalogue = install_hosted_catalogue()
|
|
2419
|
+
except Exception as error:
|
|
2420
|
+
LOGGER.warning("hosted catalogue install failed: %s", type(error).__name__)
|
|
2421
|
+
self._serve(lease, credential, deadline)
|
|
2422
|
+
except (SessionWorkerError, HostedWorkerError) as error:
|
|
2423
|
+
if getattr(error, "code", None) == "SESSION_CAPABILITY_EXPIRED":
|
|
2424
|
+
# A dispatch that reaches its absolute capability boundary during bootstrap is a
|
|
2425
|
+
# normal no-op. There is no candidate to complete or session failure to report;
|
|
2426
|
+
# Studio owns the durable close/reap and the Job must not retry it as an error.
|
|
2427
|
+
self.outcome.terminal_state = "capability_expired"
|
|
2428
|
+
self._stop.set()
|
|
2429
|
+
else:
|
|
2430
|
+
self.outcome.failure_code = getattr(error, "code", "SESSION_WORKER_FAILED")
|
|
2431
|
+
except Exception as error:
|
|
2432
|
+
LOGGER.exception("session worker failed")
|
|
2433
|
+
self.outcome.failure_code = type(error).__name__
|
|
2434
|
+
return self.outcome.to_dict()
|
|
2435
|
+
|
|
2436
|
+
def _serve(
|
|
2437
|
+
self, lease: SessionLease, credential: CredentialLease, deadline: SessionDeadline
|
|
2438
|
+
) -> None:
|
|
2439
|
+
models, _, _, _ = _load_generated()
|
|
2440
|
+
with credential.borrow() as view:
|
|
2441
|
+
client = _authenticated_client(lease.callback_base_url, bytes(view).decode("ascii"))
|
|
2442
|
+
authority = _SessionAuthority(client)
|
|
2443
|
+
sandbox = self._sandbox_factory(self._config)
|
|
2444
|
+
self.outcome.warm_pool_ready = _start_warm_pool(sandbox, self._config.warm_pool_size)
|
|
2445
|
+
transfer: Any | None = None
|
|
2446
|
+
if self._reader_factory is not None:
|
|
2447
|
+
reader = self._reader_factory(lease)
|
|
2448
|
+
else:
|
|
2449
|
+
# A SEPARATE client from the one carrying the capability. The quarantine download is
|
|
2450
|
+
# signed and must reach storage with no Authorization header on it; sharing the
|
|
2451
|
+
# authenticated client would put the attempt capability on a request to a bucket.
|
|
2452
|
+
transfer = _new_transfer_client()
|
|
2453
|
+
reader = StudioSessionAcquisitionReader(
|
|
2454
|
+
authority=authority,
|
|
2455
|
+
models=models,
|
|
2456
|
+
lease=lease,
|
|
2457
|
+
transfer=transfer,
|
|
2458
|
+
results_bucket=self._config.crawler_results_bucket,
|
|
2459
|
+
timeout_seconds=self._config.acquisition_timeout_seconds,
|
|
2460
|
+
sleep=self._sleep,
|
|
2461
|
+
)
|
|
2462
|
+
executor = ProbeExecutor(
|
|
2463
|
+
sandbox=sandbox,
|
|
2464
|
+
reader=reader,
|
|
2465
|
+
context=lease.context,
|
|
2466
|
+
)
|
|
2467
|
+
progress = _SessionProgress(
|
|
2468
|
+
authority=authority,
|
|
2469
|
+
models=models,
|
|
2470
|
+
lease=lease,
|
|
2471
|
+
deadline=deadline,
|
|
2472
|
+
)
|
|
2473
|
+
heartbeat = _Heartbeat(
|
|
2474
|
+
send=lambda: self._heartbeat(authority, models, lease, deadline),
|
|
2475
|
+
interval_seconds=lease.heartbeat_interval_seconds,
|
|
2476
|
+
stop=self._stop,
|
|
2477
|
+
deadline=deadline,
|
|
2478
|
+
on_expired=lambda: self._stop_at_deadline(deadline),
|
|
2479
|
+
)
|
|
2480
|
+
try:
|
|
2481
|
+
if self._stop_at_deadline(deadline):
|
|
2482
|
+
return
|
|
2483
|
+
# `attempt:start` is what clears `warming`. It is the first fenced callback and it is
|
|
2484
|
+
# sent before the warm pool is used for anything, so the session reports ready as soon
|
|
2485
|
+
# as a worker demonstrably exists rather than after the first probe.
|
|
2486
|
+
_response(
|
|
2487
|
+
"startProducerAttempt",
|
|
2488
|
+
authority.start_attempt.sync_detailed(
|
|
2489
|
+
attempt_id=lease.attempt_id,
|
|
2490
|
+
body=_session_callback(models, lease),
|
|
2491
|
+
idempotency_key=_idempotency("session", lease.attempt_id, "start"),
|
|
2492
|
+
),
|
|
2493
|
+
200,
|
|
2494
|
+
)
|
|
2495
|
+
if self._stop_at_deadline(deadline):
|
|
2496
|
+
return
|
|
2497
|
+
progress.emit(
|
|
2498
|
+
"attempt_started",
|
|
2499
|
+
{"role": SESSION_WORKER_ROLE, "stage_total": 0},
|
|
2500
|
+
)
|
|
2501
|
+
# Flushed here rather than left pending for the first probe. A session that opens and
|
|
2502
|
+
# then sits quiet is the ordinary case -- a person is reading before they ask anything
|
|
2503
|
+
# -- and its stream should say "the worker is here" immediately rather than at the
|
|
2504
|
+
# moment the first probe lands, which is the one moment nobody needs telling.
|
|
2505
|
+
progress.flush()
|
|
2506
|
+
if self._stop_at_deadline(deadline):
|
|
2507
|
+
return
|
|
2508
|
+
heartbeat.start()
|
|
2509
|
+
self._claim_loop(authority, models, lease, executor, progress, deadline)
|
|
2510
|
+
finally:
|
|
2511
|
+
self._stop.set()
|
|
2512
|
+
progress.emit(
|
|
2513
|
+
"attempt_finished",
|
|
2514
|
+
{"outcome": self.outcome.terminal_state or "stopped"},
|
|
2515
|
+
)
|
|
2516
|
+
progress.close()
|
|
2517
|
+
try:
|
|
2518
|
+
sandbox.close_warm_pool()
|
|
2519
|
+
except Exception:
|
|
2520
|
+
LOGGER.debug("warm pool close failed", exc_info=True)
|
|
2521
|
+
if transfer is not None:
|
|
2522
|
+
with contextlib.suppress(Exception):
|
|
2523
|
+
transfer.close()
|
|
2524
|
+
client.get_httpx_client().close()
|
|
2525
|
+
|
|
2526
|
+
def _claim_loop(
|
|
2527
|
+
self,
|
|
2528
|
+
authority: Any,
|
|
2529
|
+
models: Any,
|
|
2530
|
+
lease: SessionLease,
|
|
2531
|
+
executor: ProbeExecutor,
|
|
2532
|
+
progress: _SessionProgress,
|
|
2533
|
+
deadline: SessionDeadline,
|
|
2534
|
+
) -> None:
|
|
2535
|
+
interval = lease.claim_poll_interval_seconds or FALLBACK_CLAIM_POLL_SECONDS
|
|
2536
|
+
failures = 0
|
|
2537
|
+
poll = 0
|
|
2538
|
+
while not self._stop.is_set():
|
|
2539
|
+
if self._stop_at_deadline(deadline):
|
|
2540
|
+
return
|
|
2541
|
+
poll += 1
|
|
2542
|
+
try:
|
|
2543
|
+
answer = _response(
|
|
2544
|
+
"claimResearchProbe",
|
|
2545
|
+
authority.claim_probe.sync_detailed(
|
|
2546
|
+
attempt_id=lease.attempt_id,
|
|
2547
|
+
body=_claim_command(models, lease),
|
|
2548
|
+
# Claim is NOT idempotent-keyed on a constant. Every poll must be able to
|
|
2549
|
+
# return a different probe; a fixed key would replay the first answer for
|
|
2550
|
+
# the life of the session and the worker would answer one probe forever.
|
|
2551
|
+
# A monotonic COUNTER, not a clock reading. Two polls inside one clock
|
|
2552
|
+
# tick would share a key, and a shared claim key means Studio replaying
|
|
2553
|
+
# the first answer -- the worker would re-answer one probe while the
|
|
2554
|
+
# queue grew behind it. A counter cannot collide.
|
|
2555
|
+
idempotency_key=_idempotency(
|
|
2556
|
+
"session", lease.attempt_id, "claim", str(poll)
|
|
2557
|
+
),
|
|
2558
|
+
),
|
|
2559
|
+
200,
|
|
2560
|
+
).to_dict()
|
|
2561
|
+
except HostedWorkerError as error:
|
|
2562
|
+
if self._stop_at_deadline(deadline):
|
|
2563
|
+
return
|
|
2564
|
+
# A stale fence is Studio saying this worker is no longer the session's worker.
|
|
2565
|
+
# There is no recovery from that and retrying would be a second worker arguing
|
|
2566
|
+
# with the first.
|
|
2567
|
+
if error.code in {"ATTEMPT_FENCE_STALE", "ATTEMPT_STALE", "AUTHORIZATION_DENIED"}:
|
|
2568
|
+
self.outcome.failure_code = error.code
|
|
2569
|
+
return
|
|
2570
|
+
failures += 1
|
|
2571
|
+
if failures >= MAX_CONSECUTIVE_CLAIM_FAILURES:
|
|
2572
|
+
self.outcome.failure_code = error.code
|
|
2573
|
+
return
|
|
2574
|
+
wait_seconds = deadline.bounded_wait(interval)
|
|
2575
|
+
if wait_seconds <= 0:
|
|
2576
|
+
self._stop_at_deadline(deadline)
|
|
2577
|
+
return
|
|
2578
|
+
self._sleep(wait_seconds)
|
|
2579
|
+
continue
|
|
2580
|
+
if self._stop_at_deadline(deadline):
|
|
2581
|
+
return
|
|
2582
|
+
failures = 0
|
|
2583
|
+
state = str(answer.get("session_state", ""))
|
|
2584
|
+
interval = float(answer.get("claim_poll_interval_seconds") or interval)
|
|
2585
|
+
if state in {"closed", "failed"}:
|
|
2586
|
+
# The whole of the exit protocol. No `:complete`, no `:fail`: the session is
|
|
2587
|
+
# already terminal server-side, so both would be refused, and Studio retires the
|
|
2588
|
+
# leftover attempt row on its next sweep.
|
|
2589
|
+
self.outcome.terminal_state = state
|
|
2590
|
+
return
|
|
2591
|
+
if not answer.get("claimed"):
|
|
2592
|
+
wait_seconds = deadline.bounded_wait(interval)
|
|
2593
|
+
if wait_seconds <= 0:
|
|
2594
|
+
self._stop_at_deadline(deadline)
|
|
2595
|
+
return
|
|
2596
|
+
self._sleep(wait_seconds)
|
|
2597
|
+
continue
|
|
2598
|
+
self._answer_probe(
|
|
2599
|
+
authority,
|
|
2600
|
+
models,
|
|
2601
|
+
lease,
|
|
2602
|
+
executor,
|
|
2603
|
+
progress,
|
|
2604
|
+
deadline,
|
|
2605
|
+
answer.get("probe"),
|
|
2606
|
+
)
|
|
2607
|
+
|
|
2608
|
+
def _answer_probe(
|
|
2609
|
+
self,
|
|
2610
|
+
authority: Any,
|
|
2611
|
+
models: Any,
|
|
2612
|
+
lease: SessionLease,
|
|
2613
|
+
executor: ProbeExecutor,
|
|
2614
|
+
progress: _SessionProgress,
|
|
2615
|
+
deadline: SessionDeadline,
|
|
2616
|
+
probe: Any,
|
|
2617
|
+
) -> None:
|
|
2618
|
+
if self._stop_at_deadline(deadline):
|
|
2619
|
+
return
|
|
2620
|
+
if not isinstance(probe, dict):
|
|
2621
|
+
self.outcome.failure_code = "PROBE_CLAIM_INVALID"
|
|
2622
|
+
self._stop.set()
|
|
2623
|
+
return
|
|
2624
|
+
probe_id = str(probe.get("probe_id", ""))
|
|
2625
|
+
probe_kind = str(probe.get("probe_kind", ""))
|
|
2626
|
+
ordinal = probe.get("ordinal")
|
|
2627
|
+
ordinal = ordinal if type(ordinal) is int and ordinal > 0 else 1
|
|
2628
|
+
request = probe.get("request") or {}
|
|
2629
|
+
source_id = str(request.get("source_id", "")) if isinstance(request, dict) else ""
|
|
2630
|
+
# ONE condition for both halves of the pair. A watcher's whole use for this stream is that
|
|
2631
|
+
# a start it sees will be followed by an end, so the two records must be emittable or not
|
|
2632
|
+
# emittable together -- never a start whose end could not be spelled. `source_id` and
|
|
2633
|
+
# `probe_kind` are the two facts a settled record needs, and Studio's claim contract
|
|
2634
|
+
# requires both, so this is false only for a claim this worker could not answer anyway: a
|
|
2635
|
+
# body with no source, or a kind outside the four -- which `execute` refuses too, and which
|
|
2636
|
+
# a drift test keeps equal to the generated enum.
|
|
2637
|
+
narrated = bool(source_id) and probe_kind in session_probes.PROBE_KINDS
|
|
2638
|
+
failure_code: str | None = "PROBE_FAILED"
|
|
2639
|
+
wire_code = "PROBE_FAILED"
|
|
2640
|
+
display_code = "PROBE_FAILED"
|
|
2641
|
+
retry_after_seconds = 0
|
|
2642
|
+
result: dict[str, Any] | None = None
|
|
2643
|
+
try:
|
|
2644
|
+
if narrated:
|
|
2645
|
+
progress.emit(
|
|
2646
|
+
"source_probe_started",
|
|
2647
|
+
{"source_id": source_id, "probe_kind": probe_kind, "ordinal": ordinal},
|
|
2648
|
+
)
|
|
2649
|
+
# Flushed BEFORE the probe runs, not left to ride out with its own completion. A
|
|
2650
|
+
# probe may hold this thread for the whole `probe_timeout_seconds` ceiling, and a
|
|
2651
|
+
# start that arrives at the same instant as its end has told a watcher nothing it
|
|
2652
|
+
# did not already know.
|
|
2653
|
+
progress.flush()
|
|
2654
|
+
result = executor.execute(probe_id=probe_id, probe_kind=probe_kind, request=request)
|
|
2655
|
+
failure_code = None
|
|
2656
|
+
wire_code = progress_events.PROBE_OUTCOME_OK_CODE
|
|
2657
|
+
display_code = progress_events.PROBE_OUTCOME_OK_CODE
|
|
2658
|
+
except Exception as error:
|
|
2659
|
+
# Every exception, coded or not, and the code is KEPT when there is one. The two-clause
|
|
2660
|
+
# form this replaces named three classes and let everything else collapse to
|
|
2661
|
+
# `PROBE_FAILED` -- which silently swallowed the Clean room's own `SANDBOX_TIMEOUT`
|
|
2662
|
+
# (`AcquisitionSecurityError`, named in neither clause), turning the one failure a
|
|
2663
|
+
# client can act on into the one that says nothing.
|
|
2664
|
+
failure_code = _probe_failure_code(error)
|
|
2665
|
+
wire_code = failure_code
|
|
2666
|
+
display_code, retry_after_seconds = _probe_display_failure(error, wire_code)
|
|
2667
|
+
if failure_code == "PROBE_FAILED":
|
|
2668
|
+
LOGGER.warning("probe %s failed: %s", probe_id, type(error).__name__)
|
|
2669
|
+
if self._stop_at_deadline(deadline):
|
|
2670
|
+
return
|
|
2671
|
+
# Studio's `ContractProbeResultCommand.outcome` enum is two values. The four-value reading
|
|
2672
|
+
# below is a REFINEMENT for the unsealed timeline and never a second answer: the probe
|
|
2673
|
+
# record settles as `succeeded` or `failed`, exactly as before.
|
|
2674
|
+
settled = "succeeded" if failure_code is None else "failed"
|
|
2675
|
+
self._report(authority, models, lease, deadline, probe_id, settled, result, failure_code)
|
|
2676
|
+
if failure_code is None:
|
|
2677
|
+
self.outcome.probes_answered += 1
|
|
2678
|
+
else:
|
|
2679
|
+
self.outcome.probes_failed += 1
|
|
2680
|
+
if narrated:
|
|
2681
|
+
# EVERY terminal outcome, not only success. `source_probe_settled` travels as Studio's
|
|
2682
|
+
# `source_probe_completed` event type -- the exact type the research-session design
|
|
2683
|
+
# names for a settled probe -- so a client watching the existing per-run SSE stream
|
|
2684
|
+
# sees every probe end live, with no new event type and no second stream. Emitting
|
|
2685
|
+
# only on success (the shape this replaces) meant a watcher saw a start and never an
|
|
2686
|
+
# end, which with no source channel wired was every probe there is.
|
|
2687
|
+
#
|
|
2688
|
+
# `ordinal` and no denominator: how many probes a session will raise is not knowable
|
|
2689
|
+
# while one of them runs. The `source_acquire_*` pair this replaces demanded an
|
|
2690
|
+
# `index` and a `total`, so the worker sent the ordinal for both and every probe read
|
|
2691
|
+
# "N of N" -- a fraction a progress rail will happily draw and that never means
|
|
2692
|
+
# anything.
|
|
2693
|
+
progress.emit(
|
|
2694
|
+
"source_probe_settled",
|
|
2695
|
+
{
|
|
2696
|
+
"source_id": source_id,
|
|
2697
|
+
"probe_kind": probe_kind,
|
|
2698
|
+
"ordinal": ordinal,
|
|
2699
|
+
"outcome": _probe_progress_outcome(
|
|
2700
|
+
None if failure_code is None else display_code
|
|
2701
|
+
),
|
|
2702
|
+
"code": wire_code,
|
|
2703
|
+
"display_code": display_code,
|
|
2704
|
+
"retry_after_seconds": retry_after_seconds,
|
|
2705
|
+
"bytes": len(canonical_json_bytes(result)) if result is not None else 0,
|
|
2706
|
+
},
|
|
2707
|
+
)
|
|
2708
|
+
progress.flush()
|
|
2709
|
+
|
|
2710
|
+
def _report(
|
|
2711
|
+
self,
|
|
2712
|
+
authority: Any,
|
|
2713
|
+
models: Any,
|
|
2714
|
+
lease: SessionLease,
|
|
2715
|
+
deadline: SessionDeadline,
|
|
2716
|
+
probe_id: str,
|
|
2717
|
+
outcome: str,
|
|
2718
|
+
result: Mapping[str, Any] | None,
|
|
2719
|
+
failure_code: str | None,
|
|
2720
|
+
) -> None:
|
|
2721
|
+
if self._stop_at_deadline(deadline):
|
|
2722
|
+
return
|
|
2723
|
+
wire: dict[str, Any] = {
|
|
2724
|
+
"schema_version": START_SCHEMA_VERSION,
|
|
2725
|
+
"workspace_id": str(lease.workspace_id),
|
|
2726
|
+
"run_id": str(lease.run_id),
|
|
2727
|
+
"attempt_id": str(lease.attempt_id),
|
|
2728
|
+
"fence": lease.fence,
|
|
2729
|
+
"session_id": lease.session_id,
|
|
2730
|
+
"probe_id": probe_id,
|
|
2731
|
+
"outcome": outcome,
|
|
2732
|
+
}
|
|
2733
|
+
if result is not None:
|
|
2734
|
+
wire["result"] = dict(result)
|
|
2735
|
+
if failure_code is not None:
|
|
2736
|
+
wire["failure_code"] = failure_code
|
|
2737
|
+
try:
|
|
2738
|
+
_response(
|
|
2739
|
+
"reportResearchProbeResult",
|
|
2740
|
+
authority.report_probe_result.sync_detailed(
|
|
2741
|
+
attempt_id=lease.attempt_id,
|
|
2742
|
+
probe_id=UUID(probe_id),
|
|
2743
|
+
body=models.ContractProbeResultCommand.from_dict(wire),
|
|
2744
|
+
# Keyed on the probe, so a retried report is the same report. One answer per
|
|
2745
|
+
# probe is enforced server-side as well -- a second is
|
|
2746
|
+
# RESEARCH_PROBE_ALREADY_SETTLED -- and the key means an ordinary network
|
|
2747
|
+
# retry never becomes that.
|
|
2748
|
+
idempotency_key=_idempotency("session", lease.attempt_id, "result", probe_id),
|
|
2749
|
+
),
|
|
2750
|
+
200,
|
|
2751
|
+
)
|
|
2752
|
+
except HostedWorkerError as error:
|
|
2753
|
+
if error.code == "RESEARCH_PROBE_ALREADY_SETTLED":
|
|
2754
|
+
# Already answered, by us, before a retry. Nothing to do and nothing wrong: the
|
|
2755
|
+
# client has an answer, which is the only thing that matters.
|
|
2756
|
+
LOGGER.info("probe %s was already settled", probe_id)
|
|
2757
|
+
return
|
|
2758
|
+
raise
|
|
2759
|
+
|
|
2760
|
+
def _heartbeat(
|
|
2761
|
+
self,
|
|
2762
|
+
authority: Any,
|
|
2763
|
+
models: Any,
|
|
2764
|
+
lease: SessionLease,
|
|
2765
|
+
deadline: SessionDeadline | None = None,
|
|
2766
|
+
) -> None:
|
|
2767
|
+
if deadline is not None and self._stop_at_deadline(deadline):
|
|
2768
|
+
return
|
|
2769
|
+
_response(
|
|
2770
|
+
"heartbeatProducerAttempt",
|
|
2771
|
+
authority.heartbeat_attempt.sync_detailed(
|
|
2772
|
+
attempt_id=lease.attempt_id,
|
|
2773
|
+
body=_session_callback(models, lease),
|
|
2774
|
+
idempotency_key=_idempotency("session", lease.attempt_id, "heartbeat"),
|
|
2775
|
+
),
|
|
2776
|
+
200,
|
|
2777
|
+
)
|
|
2778
|
+
|
|
2779
|
+
|
|
2780
|
+
def _session_callback(models: Any, lease: SessionLease) -> Any:
|
|
2781
|
+
return models.ContractAttemptCallback(
|
|
2782
|
+
schema_version=START_SCHEMA_VERSION,
|
|
2783
|
+
workspace_id=lease.workspace_id,
|
|
2784
|
+
run_id=lease.run_id,
|
|
2785
|
+
attempt_id=lease.attempt_id,
|
|
2786
|
+
fence=lease.fence,
|
|
2787
|
+
)
|
|
2788
|
+
|
|
2789
|
+
|
|
2790
|
+
def _claim_command(models: Any, lease: SessionLease) -> Any:
|
|
2791
|
+
return models.ContractProbeClaimCommand.from_dict(
|
|
2792
|
+
{
|
|
2793
|
+
"schema_version": START_SCHEMA_VERSION,
|
|
2794
|
+
"workspace_id": str(lease.workspace_id),
|
|
2795
|
+
"run_id": str(lease.run_id),
|
|
2796
|
+
"attempt_id": str(lease.attempt_id),
|
|
2797
|
+
"fence": lease.fence,
|
|
2798
|
+
"session_id": lease.session_id,
|
|
2799
|
+
}
|
|
2800
|
+
)
|
|
2801
|
+
|
|
2802
|
+
|
|
2803
|
+
class _SessionProgress:
|
|
2804
|
+
"""Unsealed progress onto the session run's existing event log. Never fails a session.
|
|
2805
|
+
|
|
2806
|
+
The same promise ``hosted_worker._ProgressEmitter`` makes, kept for the same reason: a client
|
|
2807
|
+
that loses timeline resolution has lost a nicety, and a session that died because a progress
|
|
2808
|
+
append was refused has lost the thing it exists for. Records go out one at a time rather than
|
|
2809
|
+
batched, because a session's whole product claim is that a probe is visible immediately, and a
|
|
2810
|
+
250 ms coalescing window is a quarter of the follow-up budget.
|
|
2811
|
+
"""
|
|
2812
|
+
|
|
2813
|
+
def __init__(
|
|
2814
|
+
self,
|
|
2815
|
+
*,
|
|
2816
|
+
authority: Any,
|
|
2817
|
+
models: Any,
|
|
2818
|
+
lease: SessionLease,
|
|
2819
|
+
deadline: SessionDeadline | None = None,
|
|
2820
|
+
) -> None:
|
|
2821
|
+
self._authority = authority
|
|
2822
|
+
self._models = models
|
|
2823
|
+
self._lease = lease
|
|
2824
|
+
self._deadline = deadline
|
|
2825
|
+
self._sequence = 0
|
|
2826
|
+
self._batch = 0
|
|
2827
|
+
self._pending: list[Any] = []
|
|
2828
|
+
self._armed = True
|
|
2829
|
+
|
|
2830
|
+
@property
|
|
2831
|
+
def armed(self) -> bool:
|
|
2832
|
+
return self._armed
|
|
2833
|
+
|
|
2834
|
+
def emit(self, name: str, facts: Mapping[str, Any]) -> None:
|
|
2835
|
+
if self._deadline is not None and self._deadline.expired():
|
|
2836
|
+
self._armed = False
|
|
2837
|
+
self._pending.clear()
|
|
2838
|
+
return
|
|
2839
|
+
if not self._armed:
|
|
2840
|
+
return
|
|
2841
|
+
if self._sequence >= progress_events.MAX_PROGRESS_EVENTS_PER_ATTEMPT:
|
|
2842
|
+
self._armed = False
|
|
2843
|
+
self._pending.clear()
|
|
2844
|
+
return
|
|
2845
|
+
try:
|
|
2846
|
+
self._sequence += 1
|
|
2847
|
+
payload, digest = progress_events.progress_payload_and_digest(name, facts)
|
|
2848
|
+
identifier = bytearray(
|
|
2849
|
+
hashlib.sha256(
|
|
2850
|
+
f"{self._lease.attempt_id}:{self._sequence}:{name}:{digest}".encode()
|
|
2851
|
+
).digest()[:16]
|
|
2852
|
+
)
|
|
2853
|
+
identifier[6] = (identifier[6] & 0x0F) | 0x40
|
|
2854
|
+
identifier[8] = (identifier[8] & 0x3F) | 0x80
|
|
2855
|
+
self._pending.append(
|
|
2856
|
+
self._models.ContractProducerEventCommand(
|
|
2857
|
+
schema_version=START_SCHEMA_VERSION,
|
|
2858
|
+
workspace_id=self._lease.workspace_id,
|
|
2859
|
+
run_id=self._lease.run_id,
|
|
2860
|
+
attempt_id=self._lease.attempt_id,
|
|
2861
|
+
fence=self._lease.fence,
|
|
2862
|
+
client_event_id=UUID(bytes=bytes(identifier)),
|
|
2863
|
+
event_type=self._models.ContractProducerEventCommandEventType(
|
|
2864
|
+
progress_events.STUDIO_EVENT_TYPE[name]
|
|
2865
|
+
),
|
|
2866
|
+
payload_digest=digest,
|
|
2867
|
+
payload=self._models.ContractProducerEventCommandPayload.from_dict(payload),
|
|
2868
|
+
)
|
|
2869
|
+
)
|
|
2870
|
+
except Exception:
|
|
2871
|
+
LOGGER.debug("session progress record dropped", exc_info=True)
|
|
2872
|
+
|
|
2873
|
+
def flush(self) -> None:
|
|
2874
|
+
if self._deadline is not None and self._deadline.expired():
|
|
2875
|
+
self._armed = False
|
|
2876
|
+
self._pending.clear()
|
|
2877
|
+
return
|
|
2878
|
+
if not self._pending:
|
|
2879
|
+
return
|
|
2880
|
+
batch, self._pending = self._pending, []
|
|
2881
|
+
self._batch += 1
|
|
2882
|
+
if not self._armed:
|
|
2883
|
+
return
|
|
2884
|
+
try:
|
|
2885
|
+
_response(
|
|
2886
|
+
"appendProducerEvents",
|
|
2887
|
+
self._authority.append_events.sync_detailed(
|
|
2888
|
+
attempt_id=self._lease.attempt_id,
|
|
2889
|
+
body=batch,
|
|
2890
|
+
idempotency_key=_idempotency(
|
|
2891
|
+
"session", self._lease.attempt_id, "event-progress", str(self._batch)
|
|
2892
|
+
),
|
|
2893
|
+
),
|
|
2894
|
+
200,
|
|
2895
|
+
)
|
|
2896
|
+
except Exception:
|
|
2897
|
+
LOGGER.debug("session progress append failed", exc_info=True)
|
|
2898
|
+
|
|
2899
|
+
def close(self) -> None:
|
|
2900
|
+
try:
|
|
2901
|
+
self.flush()
|
|
2902
|
+
finally:
|
|
2903
|
+
self._armed = False
|
|
2904
|
+
|
|
2905
|
+
|
|
2906
|
+
def _start_warm_pool(sandbox: Any, size: int) -> int:
|
|
2907
|
+
"""Ask for a warm pool and believe the answer.
|
|
2908
|
+
|
|
2909
|
+
``start_warm_pool`` returns the number now ready, which on Linux under a bounded cgroup root is
|
|
2910
|
+
zero by design and is not an error -- the first probe simply spawns cold. The lazy-import gain
|
|
2911
|
+
that came with the pool (#306) applies either way, so a zero costs latency and nothing else.
|
|
2912
|
+
"""
|
|
2913
|
+
|
|
2914
|
+
if size <= 0:
|
|
2915
|
+
return 0
|
|
2916
|
+
try:
|
|
2917
|
+
return int(sandbox.start_warm_pool(size=size))
|
|
2918
|
+
except Exception:
|
|
2919
|
+
LOGGER.info("warm sandbox pool unavailable; probes will spawn cold")
|
|
2920
|
+
return 0
|
|
2921
|
+
|
|
2922
|
+
|
|
2923
|
+
def _default_sandbox(config: SessionWorkerConfig) -> Any:
|
|
2924
|
+
from mostlyright.data_harness.acquisition.sandbox import (
|
|
2925
|
+
HOSTED_SESSION_CONTAINER_MEMORY_BYTES,
|
|
2926
|
+
CrawlerSandbox,
|
|
2927
|
+
)
|
|
2928
|
+
|
|
2929
|
+
staging_root = Path(tempfile.mkdtemp(prefix="mr-session-probe-")).resolve()
|
|
2930
|
+
return CrawlerSandbox(
|
|
2931
|
+
staging_root=staging_root,
|
|
2932
|
+
timeout_seconds=config.probe_timeout_seconds,
|
|
2933
|
+
external_network_policy_attestation=config.external_network_policy_attestation,
|
|
2934
|
+
hosted_session_container_memory_bytes=HOSTED_SESSION_CONTAINER_MEMORY_BYTES,
|
|
2935
|
+
)
|
|
2936
|
+
|
|
2937
|
+
|
|
2938
|
+
# ==========================================================================================
|
|
2939
|
+
# The HTTP surface
|
|
2940
|
+
# ==========================================================================================
|
|
2941
|
+
|
|
2942
|
+
|
|
2943
|
+
class _LegacySessionProcess:
|
|
2944
|
+
"""A capability-free Service-side handle for one fresh legacy-session interpreter.
|
|
2945
|
+
|
|
2946
|
+
``ThreadingHTTPServer`` calls :meth:`SessionWorkerService.start` from a request thread. It is
|
|
2947
|
+
unsafe to ``fork()`` there and then execute Python, HTTP clients or parsers: another thread
|
|
2948
|
+
may have held one of their runtime locks at the instant of the fork. This handle uses the
|
|
2949
|
+
platform's subprocess exec path only, then the fresh interpreter creates the monitored
|
|
2950
|
+
capability child from its main thread. The Service process therefore retains neither the
|
|
2951
|
+
lease capability nor a post-fork Python child.
|
|
2952
|
+
"""
|
|
2953
|
+
|
|
2954
|
+
def __init__(self, *, process: Any, session_id: str) -> None:
|
|
2955
|
+
self._process = process
|
|
2956
|
+
self._session_id = session_id
|
|
2957
|
+
self._stopped = False
|
|
2958
|
+
|
|
2959
|
+
@classmethod
|
|
2960
|
+
def start(
|
|
2961
|
+
cls,
|
|
2962
|
+
*,
|
|
2963
|
+
config: SessionWorkerConfig,
|
|
2964
|
+
invocation: StartInvocation,
|
|
2965
|
+
nonce: bytearray,
|
|
2966
|
+
) -> _LegacySessionProcess:
|
|
2967
|
+
"""Pass one validated nonce over a pipe to a fresh interpreter, then zero local copies."""
|
|
2968
|
+
|
|
2969
|
+
try:
|
|
2970
|
+
read_descriptor, write_descriptor = os.pipe()
|
|
2971
|
+
except OSError as error:
|
|
2972
|
+
nonce[:] = b"\x00" * len(nonce)
|
|
2973
|
+
raise SessionWorkerError(
|
|
2974
|
+
"SESSION_LEGACY_PROCESS_START_FAILED",
|
|
2975
|
+
"could not allocate the isolated legacy session pipe",
|
|
2976
|
+
) from error
|
|
2977
|
+
process: Any | None = None
|
|
2978
|
+
payload = bytearray()
|
|
2979
|
+
try:
|
|
2980
|
+
# No ``preexec_fn``: on POSIX that would run Python after a multithreaded fork. The
|
|
2981
|
+
# executable receives only a descriptor number; the nonce is written after exec over
|
|
2982
|
+
# an inherited private pipe and never appears in argv or the environment.
|
|
2983
|
+
process = subprocess.Popen(
|
|
2984
|
+
(
|
|
2985
|
+
sys.executable,
|
|
2986
|
+
"-m",
|
|
2987
|
+
"mostlyright.data_harness.hosted_session_worker",
|
|
2988
|
+
LEGACY_BOOTSTRAP_FD_ARGUMENT,
|
|
2989
|
+
str(read_descriptor),
|
|
2990
|
+
),
|
|
2991
|
+
close_fds=True,
|
|
2992
|
+
pass_fds=(read_descriptor,),
|
|
2993
|
+
stdin=subprocess.DEVNULL,
|
|
2994
|
+
stdout=subprocess.PIPE,
|
|
2995
|
+
start_new_session=True,
|
|
2996
|
+
)
|
|
2997
|
+
os.close(read_descriptor)
|
|
2998
|
+
read_descriptor = -1
|
|
2999
|
+
payload = bytearray(
|
|
3000
|
+
canonical_json_bytes(
|
|
3001
|
+
{
|
|
3002
|
+
"schema_version": START_SCHEMA_VERSION,
|
|
3003
|
+
"session_id": invocation.session_id,
|
|
3004
|
+
"workspace_id": invocation.workspace_id,
|
|
3005
|
+
"lease_nonce": nonce.decode("ascii"),
|
|
3006
|
+
"lease_url": invocation.lease_url,
|
|
3007
|
+
}
|
|
3008
|
+
)
|
|
3009
|
+
)
|
|
3010
|
+
if len(payload) > MAX_LEGACY_BOOTSTRAP_BYTES:
|
|
3011
|
+
raise SessionWorkerError(
|
|
3012
|
+
"SESSION_LEGACY_BOOTSTRAP_INVALID",
|
|
3013
|
+
"legacy bootstrap exceeds its byte limit",
|
|
3014
|
+
)
|
|
3015
|
+
cls._write_all(write_descriptor, payload)
|
|
3016
|
+
cls._await_ready(process)
|
|
3017
|
+
except (OSError, UnicodeError, SessionWorkerError, subprocess.SubprocessError) as error:
|
|
3018
|
+
if process is not None:
|
|
3019
|
+
cls._terminate_process_group(process)
|
|
3020
|
+
raise SessionWorkerError(
|
|
3021
|
+
"SESSION_LEGACY_PROCESS_START_FAILED",
|
|
3022
|
+
"could not start the isolated legacy session process",
|
|
3023
|
+
) from error
|
|
3024
|
+
finally:
|
|
3025
|
+
if read_descriptor >= 0:
|
|
3026
|
+
with contextlib.suppress(OSError):
|
|
3027
|
+
os.close(read_descriptor)
|
|
3028
|
+
with contextlib.suppress(OSError):
|
|
3029
|
+
os.close(write_descriptor)
|
|
3030
|
+
payload[:] = b"\x00" * len(payload)
|
|
3031
|
+
nonce[:] = b"\x00" * len(nonce)
|
|
3032
|
+
assert process is not None
|
|
3033
|
+
return cls(process=process, session_id=invocation.session_id)
|
|
3034
|
+
|
|
3035
|
+
@staticmethod
|
|
3036
|
+
def _write_all(descriptor: int, payload: bytearray) -> None:
|
|
3037
|
+
view = memoryview(payload)
|
|
3038
|
+
try:
|
|
3039
|
+
while view:
|
|
3040
|
+
written = os.write(descriptor, view)
|
|
3041
|
+
if written <= 0:
|
|
3042
|
+
raise OSError("short legacy bootstrap pipe write")
|
|
3043
|
+
view = view[written:]
|
|
3044
|
+
finally:
|
|
3045
|
+
view.release()
|
|
3046
|
+
|
|
3047
|
+
@staticmethod
|
|
3048
|
+
def _await_ready(process: Any) -> None:
|
|
3049
|
+
"""Require the fresh bootstrap monitor to arm its cleanup handler before HTTP accepts."""
|
|
3050
|
+
|
|
3051
|
+
stdout = process.stdout
|
|
3052
|
+
if stdout is None:
|
|
3053
|
+
raise OSError("legacy bootstrap process has no control output")
|
|
3054
|
+
readable, _unused, _errors = select.select(
|
|
3055
|
+
[stdout.fileno()], [], [], LEGACY_BOOTSTRAP_READY_TIMEOUT_SECONDS
|
|
3056
|
+
)
|
|
3057
|
+
if not readable:
|
|
3058
|
+
raise OSError("legacy bootstrap monitor did not become ready in time")
|
|
3059
|
+
raw = stdout.readline(MAX_SESSION_SUPERVISOR_MESSAGE_BYTES + 1)
|
|
3060
|
+
if len(raw) > MAX_SESSION_SUPERVISOR_MESSAGE_BYTES or not raw.endswith(b"\n"):
|
|
3061
|
+
raise OSError("legacy bootstrap monitor readiness is invalid")
|
|
3062
|
+
try:
|
|
3063
|
+
message = _strict_json(raw[:-1], "legacy bootstrap readiness")
|
|
3064
|
+
except HostedWorkerError as error:
|
|
3065
|
+
raise OSError("legacy bootstrap monitor readiness is invalid") from error
|
|
3066
|
+
if message != {"kind": "ready"}:
|
|
3067
|
+
raise OSError("legacy bootstrap monitor did not acknowledge readiness")
|
|
3068
|
+
|
|
3069
|
+
@staticmethod
|
|
3070
|
+
def _terminate_process_group(process: Any) -> None:
|
|
3071
|
+
if process.poll() is not None:
|
|
3072
|
+
return
|
|
3073
|
+
# The fresh bootstrap installs its monitor handler *before* it forks the separate
|
|
3074
|
+
# capability group. SIGTERM therefore lets it kill/reap that group and its adopted
|
|
3075
|
+
# sandboxes during a readiness failure; SIGKILL is only the bounded fallback for a wedged
|
|
3076
|
+
# outer monitor.
|
|
3077
|
+
with contextlib.suppress(PermissionError, ProcessLookupError):
|
|
3078
|
+
os.killpg(process.pid, signal.SIGTERM)
|
|
3079
|
+
try:
|
|
3080
|
+
process.wait(timeout=1.0)
|
|
3081
|
+
return
|
|
3082
|
+
except subprocess.TimeoutExpired:
|
|
3083
|
+
pass
|
|
3084
|
+
with contextlib.suppress(PermissionError, ProcessLookupError):
|
|
3085
|
+
os.killpg(process.pid, signal.SIGKILL)
|
|
3086
|
+
with contextlib.suppress(Exception):
|
|
3087
|
+
process.wait(timeout=5)
|
|
3088
|
+
|
|
3089
|
+
@property
|
|
3090
|
+
def is_running(self) -> bool:
|
|
3091
|
+
return self._process.poll() is None
|
|
3092
|
+
|
|
3093
|
+
def stop(self) -> None:
|
|
3094
|
+
self._stopped = True
|
|
3095
|
+
if self._process.poll() is not None:
|
|
3096
|
+
return
|
|
3097
|
+
# The fresh bootstrap interpreter catches this signal in its capability-free monitor, then
|
|
3098
|
+
# kills the separate capability group and reaps its adopted sandbox groups before exiting.
|
|
3099
|
+
# The group here intentionally contains only that bootstrap interpreter, not its child.
|
|
3100
|
+
with contextlib.suppress(ProcessLookupError):
|
|
3101
|
+
os.killpg(self._process.pid, signal.SIGTERM)
|
|
3102
|
+
|
|
3103
|
+
def run(self) -> dict[str, Any]:
|
|
3104
|
+
raw, _unused_stderr = self._process.communicate()
|
|
3105
|
+
if (
|
|
3106
|
+
self._stopped
|
|
3107
|
+
and not raw
|
|
3108
|
+
and self._process.returncode is not None
|
|
3109
|
+
and self._process.returncode < 0
|
|
3110
|
+
):
|
|
3111
|
+
return {
|
|
3112
|
+
"status": "stopped",
|
|
3113
|
+
"role": SESSION_WORKER_ROLE,
|
|
3114
|
+
"session_id": self._session_id,
|
|
3115
|
+
"attempt_id": None,
|
|
3116
|
+
"probes_answered": 0,
|
|
3117
|
+
"probes_failed": 0,
|
|
3118
|
+
"terminal_state": "service_stopped",
|
|
3119
|
+
"failure_code": None,
|
|
3120
|
+
"warm_pool_ready": 0,
|
|
3121
|
+
}
|
|
3122
|
+
if (
|
|
3123
|
+
self._process.returncode == 0
|
|
3124
|
+
and len(raw) <= MAX_SESSION_SUPERVISOR_MESSAGE_BYTES
|
|
3125
|
+
and raw.endswith(b"\n")
|
|
3126
|
+
and raw.count(b"\n") == 1
|
|
3127
|
+
):
|
|
3128
|
+
try:
|
|
3129
|
+
summary = _strict_json(raw[:-1], "isolated legacy session summary")
|
|
3130
|
+
except HostedWorkerError:
|
|
3131
|
+
summary = None
|
|
3132
|
+
if isinstance(summary, dict) and summary.get("session_id") == self._session_id:
|
|
3133
|
+
return summary
|
|
3134
|
+
return {
|
|
3135
|
+
"status": "failed",
|
|
3136
|
+
"role": SESSION_WORKER_ROLE,
|
|
3137
|
+
"session_id": self._session_id,
|
|
3138
|
+
"attempt_id": None,
|
|
3139
|
+
"probes_answered": 0,
|
|
3140
|
+
"probes_failed": 0,
|
|
3141
|
+
"terminal_state": None,
|
|
3142
|
+
"failure_code": "SESSION_LEGACY_PROCESS_FAILED",
|
|
3143
|
+
"warm_pool_ready": 0,
|
|
3144
|
+
}
|
|
3145
|
+
|
|
3146
|
+
|
|
3147
|
+
class SessionWorkerService:
|
|
3148
|
+
"""One process, at most one live session, and an exit that lets Cloud Run scale to zero.
|
|
3149
|
+
|
|
3150
|
+
``max_instance_request_concurrency = 1`` means Cloud Run will not send this instance a second
|
|
3151
|
+
request while the start invocation is being handled, but the invocation returns in
|
|
3152
|
+
milliseconds, so a second start *can* arrive while the first session is still running. It is
|
|
3153
|
+
refused with 409 rather than served: one session per instance is the tofu's own comment ("A
|
|
3154
|
+
probe must not queue behind another session's CPU") and serving two here would silently break
|
|
3155
|
+
the latency claim the whole feature is built on. Cloud Run answers a 409 by placing the session
|
|
3156
|
+
on another instance, which is what should happen.
|
|
3157
|
+
|
|
3158
|
+
Sessions are served serially, though, and that matters: a client who closes one session and
|
|
3159
|
+
opens another within the grace window gets this warm instance instead of a second cold start.
|
|
3160
|
+
"""
|
|
3161
|
+
|
|
3162
|
+
def __init__(
|
|
3163
|
+
self,
|
|
3164
|
+
config: SessionWorkerConfig,
|
|
3165
|
+
*,
|
|
3166
|
+
supervisor_factory: Callable[..., SessionSupervisor] | None = None,
|
|
3167
|
+
) -> None:
|
|
3168
|
+
self._config = config
|
|
3169
|
+
self._supervisor_factory = supervisor_factory
|
|
3170
|
+
self._lock = threading.Lock()
|
|
3171
|
+
self._active: Any | None = None
|
|
3172
|
+
self._idle_since: float | None = time.monotonic()
|
|
3173
|
+
self._shutdown = threading.Event()
|
|
3174
|
+
self.completed: list[dict[str, Any]] = []
|
|
3175
|
+
|
|
3176
|
+
# -- the two verbs ----------------------------------------------------------------------
|
|
3177
|
+
|
|
3178
|
+
def start(self, raw: bytes) -> tuple[int, dict[str, Any]]:
|
|
3179
|
+
"""Accept one start signal and return in milliseconds. See the module docstring."""
|
|
3180
|
+
|
|
3181
|
+
try:
|
|
3182
|
+
invocation, nonce = parse_start_invocation(
|
|
3183
|
+
raw, studio_origin=self._config.studio_origin
|
|
3184
|
+
)
|
|
3185
|
+
except SessionWorkerError as error:
|
|
3186
|
+
# The reason is returned because Studio's starter treats any non-2xx as a failed
|
|
3187
|
+
# session and an operator reading the log needs to know which check refused. The nonce
|
|
3188
|
+
# is not in the invocation object, so nothing here can echo it.
|
|
3189
|
+
return int(HTTPStatus.BAD_REQUEST), {"code": error.code, "detail": error.detail}
|
|
3190
|
+
with self._lock:
|
|
3191
|
+
if self._active is not None and self._active.is_running:
|
|
3192
|
+
nonce[:] = b"\x00" * len(nonce)
|
|
3193
|
+
return int(HTTPStatus.CONFLICT), {
|
|
3194
|
+
"code": "SESSION_WORKER_BUSY",
|
|
3195
|
+
"detail": "this instance already holds a live session",
|
|
3196
|
+
}
|
|
3197
|
+
if self._supervisor_factory is None:
|
|
3198
|
+
try:
|
|
3199
|
+
supervisor: Any = _LegacySessionProcess.start(
|
|
3200
|
+
config=self._config, invocation=invocation, nonce=nonce
|
|
3201
|
+
)
|
|
3202
|
+
except SessionWorkerError as error:
|
|
3203
|
+
return int(HTTPStatus.SERVICE_UNAVAILABLE), {
|
|
3204
|
+
"code": error.code,
|
|
3205
|
+
"detail": error.detail,
|
|
3206
|
+
}
|
|
3207
|
+
else:
|
|
3208
|
+
supervisor = self._supervisor_factory(
|
|
3209
|
+
config=self._config, invocation=invocation, nonce=nonce
|
|
3210
|
+
)
|
|
3211
|
+
self._active = supervisor
|
|
3212
|
+
self._idle_since = None
|
|
3213
|
+
threading.Thread(
|
|
3214
|
+
target=self._run_session,
|
|
3215
|
+
args=(supervisor,),
|
|
3216
|
+
name=f"mr-session-{invocation.session_id[:8]}",
|
|
3217
|
+
daemon=False,
|
|
3218
|
+
).start()
|
|
3219
|
+
return int(HTTPStatus.ACCEPTED), {
|
|
3220
|
+
"accepted": True,
|
|
3221
|
+
"session_id": invocation.session_id,
|
|
3222
|
+
}
|
|
3223
|
+
|
|
3224
|
+
def live(self) -> tuple[int, dict[str, Any]]:
|
|
3225
|
+
"""Liveness. Deliberately says nothing about whether a session is warm."""
|
|
3226
|
+
|
|
3227
|
+
return int(HTTPStatus.OK), {"status": "live", "role": SESSION_WORKER_ROLE}
|
|
3228
|
+
|
|
3229
|
+
# -- lifecycle --------------------------------------------------------------------------
|
|
3230
|
+
|
|
3231
|
+
def _run_session(self, supervisor: SessionSupervisor) -> None:
|
|
3232
|
+
try:
|
|
3233
|
+
summary = supervisor.run()
|
|
3234
|
+
finally:
|
|
3235
|
+
with self._lock:
|
|
3236
|
+
self._active = None
|
|
3237
|
+
self._idle_since = time.monotonic()
|
|
3238
|
+
self.completed.append(summary)
|
|
3239
|
+
sys.stdout.write(json.dumps(summary, sort_keys=True, separators=(",", ":")) + "\n")
|
|
3240
|
+
sys.stdout.flush()
|
|
3241
|
+
|
|
3242
|
+
def idle_seconds(self, *, now: float | None = None) -> float | None:
|
|
3243
|
+
with self._lock:
|
|
3244
|
+
if self._idle_since is None:
|
|
3245
|
+
return None
|
|
3246
|
+
return (time.monotonic() if now is None else now) - self._idle_since
|
|
3247
|
+
|
|
3248
|
+
def should_exit(self, *, now: float | None = None) -> bool:
|
|
3249
|
+
"""True once the instance has been idle long enough to be worth giving back.
|
|
3250
|
+
|
|
3251
|
+
This is the whole of scale-to-zero on the worker side. Cloud Run would eventually reap an
|
|
3252
|
+
idle instance on its own, but 'eventually' is billed at ``cpu_idle = false`` rates, and the
|
|
3253
|
+
difference between roughly $21 and roughly $55 per user-month is the difference between the
|
|
3254
|
+
product's price working and not.
|
|
3255
|
+
"""
|
|
3256
|
+
|
|
3257
|
+
idle = self.idle_seconds(now=now)
|
|
3258
|
+
return idle is not None and idle >= self._config.shutdown_grace_seconds
|
|
3259
|
+
|
|
3260
|
+
def stop(self) -> None:
|
|
3261
|
+
with self._lock:
|
|
3262
|
+
active = self._active
|
|
3263
|
+
if active is not None:
|
|
3264
|
+
active.stop()
|
|
3265
|
+
self._shutdown.set()
|
|
3266
|
+
|
|
3267
|
+
|
|
3268
|
+
class _Handler(BaseHTTPRequestHandler):
|
|
3269
|
+
"""The smallest surface that satisfies the contract. Everything else is 404."""
|
|
3270
|
+
|
|
3271
|
+
protocol_version = "HTTP/1.1"
|
|
3272
|
+
service: SessionWorkerService
|
|
3273
|
+
|
|
3274
|
+
def version_string(self) -> str:
|
|
3275
|
+
return "mostlyright-session-worker"
|
|
3276
|
+
|
|
3277
|
+
def log_message(self, format: str, *args: Any) -> None:
|
|
3278
|
+
# The default writes the full request line to stderr. A request line is never secret here
|
|
3279
|
+
# (the nonce is in the body), but routing it through the logger keeps one place responsible
|
|
3280
|
+
# for what this service says out loud.
|
|
3281
|
+
LOGGER.info("%s %s", self.command, self.path)
|
|
3282
|
+
|
|
3283
|
+
def do_GET(self) -> None:
|
|
3284
|
+
if self.path in HEALTH_PATHS:
|
|
3285
|
+
self._respond(*self.service.live())
|
|
3286
|
+
return
|
|
3287
|
+
self._respond(int(HTTPStatus.NOT_FOUND), {"code": "NOT_FOUND"})
|
|
3288
|
+
|
|
3289
|
+
def do_POST(self) -> None:
|
|
3290
|
+
if self.path != START_PATH:
|
|
3291
|
+
self._respond(int(HTTPStatus.NOT_FOUND), {"code": "NOT_FOUND"})
|
|
3292
|
+
return
|
|
3293
|
+
try:
|
|
3294
|
+
length = int(self.headers.get("Content-Length", "0"))
|
|
3295
|
+
except ValueError:
|
|
3296
|
+
self._respond(int(HTTPStatus.BAD_REQUEST), {"code": "SESSION_START_INVALID"})
|
|
3297
|
+
return
|
|
3298
|
+
if length < 0 or length > MAX_START_REQUEST_BYTES:
|
|
3299
|
+
self._respond(
|
|
3300
|
+
int(HTTPStatus.REQUEST_ENTITY_TOO_LARGE), {"code": "SESSION_START_TOO_LARGE"}
|
|
3301
|
+
)
|
|
3302
|
+
return
|
|
3303
|
+
raw = self.rfile.read(length)
|
|
3304
|
+
self._respond(*self.service.start(raw))
|
|
3305
|
+
|
|
3306
|
+
def _respond(self, status: int, body: Mapping[str, Any]) -> None:
|
|
3307
|
+
encoded = json.dumps(body, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
|
3308
|
+
self.send_response(status)
|
|
3309
|
+
self.send_header("Content-Type", "application/json")
|
|
3310
|
+
self.send_header("Content-Length", str(len(encoded)))
|
|
3311
|
+
self.send_header("Cache-Control", "no-store")
|
|
3312
|
+
self.end_headers()
|
|
3313
|
+
self.wfile.write(encoded)
|
|
3314
|
+
|
|
3315
|
+
|
|
3316
|
+
class _Server(ThreadingHTTPServer):
|
|
3317
|
+
daemon_threads = True
|
|
3318
|
+
allow_reuse_address = True
|
|
3319
|
+
|
|
3320
|
+
|
|
3321
|
+
def serve(
|
|
3322
|
+
config: SessionWorkerConfig,
|
|
3323
|
+
*,
|
|
3324
|
+
service: SessionWorkerService | None = None,
|
|
3325
|
+
poll_seconds: float = 0.5,
|
|
3326
|
+
) -> int:
|
|
3327
|
+
"""Run the service until the instance has been idle long enough to give back.
|
|
3328
|
+
|
|
3329
|
+
The exit is the point. A container that stays up holds a Cloud Run instance, and an instance
|
|
3330
|
+
that is held is billed; returning from here lets the platform drop it immediately rather than
|
|
3331
|
+
waiting on its own idle reaper.
|
|
3332
|
+
"""
|
|
3333
|
+
|
|
3334
|
+
worker = service or SessionWorkerService(config)
|
|
3335
|
+
handler = type("_BoundHandler", (_Handler,), {"service": worker})
|
|
3336
|
+
shutdown_requested = threading.Event()
|
|
3337
|
+
previous_sigterm: Any | None = None
|
|
3338
|
+
if threading.current_thread() is threading.main_thread():
|
|
3339
|
+
previous_sigterm = signal.getsignal(signal.SIGTERM)
|
|
3340
|
+
|
|
3341
|
+
def request_shutdown(_signum: int, _frame: Any) -> None:
|
|
3342
|
+
# A signal handler must not take the Service lock. The main lifecycle loop calls
|
|
3343
|
+
# ``worker.stop`` and lets its bootstrap monitor kill/reap the isolated child.
|
|
3344
|
+
shutdown_requested.set()
|
|
3345
|
+
|
|
3346
|
+
signal.signal(signal.SIGTERM, request_shutdown)
|
|
3347
|
+
with _Server(("0.0.0.0", config.port), handler) as server:
|
|
3348
|
+
thread = threading.Thread(target=server.serve_forever, name="mr-session-http", daemon=True)
|
|
3349
|
+
thread.start()
|
|
3350
|
+
try:
|
|
3351
|
+
while not shutdown_requested.is_set() and not worker.should_exit():
|
|
3352
|
+
time.sleep(poll_seconds)
|
|
3353
|
+
finally:
|
|
3354
|
+
worker.stop()
|
|
3355
|
+
server.shutdown()
|
|
3356
|
+
if previous_sigterm is not None:
|
|
3357
|
+
signal.signal(signal.SIGTERM, previous_sigterm)
|
|
3358
|
+
return 0
|
|
3359
|
+
|
|
3360
|
+
|
|
3361
|
+
def run_session_job(
|
|
3362
|
+
*,
|
|
3363
|
+
config: SessionWorkerConfig,
|
|
3364
|
+
dispatch: SessionDispatch,
|
|
3365
|
+
transport: BootstrapTransport | None = None,
|
|
3366
|
+
sandbox_factory: Callable[[SessionWorkerConfig], Any] | None = None,
|
|
3367
|
+
reader_factory: Callable[[SessionLease], ProbeSourceReader] | None = None,
|
|
3368
|
+
sleep: Callable[[float], None] = time.sleep,
|
|
3369
|
+
wall_clock: Callable[[], datetime] = _wall_clock_now,
|
|
3370
|
+
monotonic_clock: Callable[[], float] = time.monotonic,
|
|
3371
|
+
process_isolation: bool = True,
|
|
3372
|
+
) -> dict[str, Any]:
|
|
3373
|
+
"""Exchange and synchronously supervise exactly one session in one Cloud Run Job task."""
|
|
3374
|
+
|
|
3375
|
+
# The command-line audience is the task's lease audience. Do not read a legacy audience
|
|
3376
|
+
# environment variable after the dispatch has been accepted: that would make a task's OIDC
|
|
3377
|
+
# proof differ from the deployment coordinate Studio minted for it.
|
|
3378
|
+
job_config = replace(config, lease_audience=dispatch.audience)
|
|
3379
|
+
if process_isolation:
|
|
3380
|
+
# The opaque dispatch is safe in the parent; the capability exchange is intentionally in
|
|
3381
|
+
# the child, so neither this monitor nor the Cloud Run task's launcher retain a bearer.
|
|
3382
|
+
supervisor = SessionSupervisor(
|
|
3383
|
+
config=job_config,
|
|
3384
|
+
dispatch=dispatch,
|
|
3385
|
+
transport=transport,
|
|
3386
|
+
sandbox_factory=sandbox_factory,
|
|
3387
|
+
reader_factory=reader_factory,
|
|
3388
|
+
sleep=sleep,
|
|
3389
|
+
wall_clock=wall_clock,
|
|
3390
|
+
monotonic_clock=monotonic_clock,
|
|
3391
|
+
process_isolation=True,
|
|
3392
|
+
)
|
|
3393
|
+
else:
|
|
3394
|
+
# Deterministic unit adapters are intentionally in-process; production callers use the
|
|
3395
|
+
# default above. This seam prevents forked copies of in-memory fakes from being mistaken
|
|
3396
|
+
# for an end-to-end Studio result.
|
|
3397
|
+
lease, credential = exchange_session_dispatch(
|
|
3398
|
+
config=job_config, dispatch=dispatch, transport=transport
|
|
3399
|
+
)
|
|
3400
|
+
supervisor = SessionSupervisor(
|
|
3401
|
+
config=job_config,
|
|
3402
|
+
bootstrap=(lease, credential),
|
|
3403
|
+
transport=transport,
|
|
3404
|
+
sandbox_factory=sandbox_factory,
|
|
3405
|
+
reader_factory=reader_factory,
|
|
3406
|
+
sleep=sleep,
|
|
3407
|
+
wall_clock=wall_clock,
|
|
3408
|
+
monotonic_clock=monotonic_clock,
|
|
3409
|
+
)
|
|
3410
|
+
return supervisor.run()
|
|
3411
|
+
|
|
3412
|
+
|
|
3413
|
+
def _read_legacy_bootstrap_descriptor(descriptor: int) -> bytearray:
|
|
3414
|
+
"""Read exactly one bounded start request from the Service's private inherited pipe."""
|
|
3415
|
+
|
|
3416
|
+
if descriptor < 3:
|
|
3417
|
+
raise SessionWorkerError(
|
|
3418
|
+
"SESSION_LEGACY_BOOTSTRAP_INVALID", "legacy bootstrap descriptor is not inherited"
|
|
3419
|
+
)
|
|
3420
|
+
raw = bytearray()
|
|
3421
|
+
try:
|
|
3422
|
+
while True:
|
|
3423
|
+
chunk = os.read(descriptor, min(4096, MAX_LEGACY_BOOTSTRAP_BYTES + 1 - len(raw)))
|
|
3424
|
+
if not chunk:
|
|
3425
|
+
break
|
|
3426
|
+
raw.extend(chunk)
|
|
3427
|
+
if len(raw) > MAX_LEGACY_BOOTSTRAP_BYTES:
|
|
3428
|
+
raise SessionWorkerError(
|
|
3429
|
+
"SESSION_LEGACY_BOOTSTRAP_INVALID", "legacy bootstrap exceeds its byte limit"
|
|
3430
|
+
)
|
|
3431
|
+
finally:
|
|
3432
|
+
with contextlib.suppress(OSError):
|
|
3433
|
+
os.close(descriptor)
|
|
3434
|
+
if not raw:
|
|
3435
|
+
raise SessionWorkerError("SESSION_LEGACY_BOOTSTRAP_INVALID", "legacy bootstrap is empty")
|
|
3436
|
+
return raw
|
|
3437
|
+
|
|
3438
|
+
|
|
3439
|
+
def _legacy_bootstrap_child_main(arguments: Sequence[str]) -> int:
|
|
3440
|
+
"""Start the rollback-session supervisor in a fresh main-thread interpreter.
|
|
3441
|
+
|
|
3442
|
+
This is intentionally an internal descriptor protocol rather than an alternate public Service
|
|
3443
|
+
command. It performs the old nonce lease exchange only inside a newly exec'd process, whose
|
|
3444
|
+
own :class:`SessionSupervisor` then forks the capability-bearing child from a safe main
|
|
3445
|
+
thread. Neither parent can be held past capability expiry by a reader or native parser.
|
|
3446
|
+
"""
|
|
3447
|
+
|
|
3448
|
+
if (
|
|
3449
|
+
len(arguments) != 2
|
|
3450
|
+
or arguments[0] != LEGACY_BOOTSTRAP_FD_ARGUMENT
|
|
3451
|
+
or not arguments[1].isdecimal()
|
|
3452
|
+
):
|
|
3453
|
+
sys.stderr.write("SESSION_LEGACY_BOOTSTRAP_INVALID: expected one inherited descriptor\n")
|
|
3454
|
+
return 2
|
|
3455
|
+
try:
|
|
3456
|
+
descriptor = int(arguments[1])
|
|
3457
|
+
raw = _read_legacy_bootstrap_descriptor(descriptor)
|
|
3458
|
+
config = config_from_environment()
|
|
3459
|
+
invocation, nonce = parse_start_invocation(raw, studio_origin=config.studio_origin)
|
|
3460
|
+
except (SessionWorkerError, HostedWorkerError, OSError) as error:
|
|
3461
|
+
code = getattr(error, "code", "SESSION_LEGACY_BOOTSTRAP_INVALID")
|
|
3462
|
+
sys.stderr.write(f"{code}: bootstrap refused\n")
|
|
3463
|
+
return 1
|
|
3464
|
+
finally:
|
|
3465
|
+
if "raw" in locals():
|
|
3466
|
+
raw[:] = b"\x00" * len(raw)
|
|
3467
|
+
|
|
3468
|
+
try:
|
|
3469
|
+
summary = SessionSupervisor(
|
|
3470
|
+
config=config,
|
|
3471
|
+
invocation=invocation,
|
|
3472
|
+
nonce=nonce,
|
|
3473
|
+
process_isolation=True,
|
|
3474
|
+
on_isolation_ready=lambda: SessionSupervisor._write_isolated_message(
|
|
3475
|
+
sys.stdout.fileno(), {"kind": "ready"}
|
|
3476
|
+
),
|
|
3477
|
+
).run()
|
|
3478
|
+
finally:
|
|
3479
|
+
nonce[:] = b"\x00" * len(nonce)
|
|
3480
|
+
try:
|
|
3481
|
+
SessionSupervisor._write_isolated_message(sys.stdout.fileno(), summary)
|
|
3482
|
+
except (OSError, SessionWorkerError):
|
|
3483
|
+
return 1
|
|
3484
|
+
return 0 if summary["status"] == "stopped" else 1
|
|
3485
|
+
|
|
3486
|
+
|
|
3487
|
+
def session_worker_job_main(argv: Sequence[str] | None = None) -> int:
|
|
3488
|
+
"""The Cloud Run Job ``mr-data-session-worker-job`` entrypoint.
|
|
3489
|
+
|
|
3490
|
+
Its sole task-specific inputs are the opaque dispatch id, exact exchange URL and OIDC audience;
|
|
3491
|
+
the task never receives a nonce or a capability in argv or its environment.
|
|
3492
|
+
"""
|
|
3493
|
+
|
|
3494
|
+
logging.basicConfig(
|
|
3495
|
+
level=logging.INFO, stream=sys.stderr, format="%(levelname)s %(name)s %(message)s"
|
|
3496
|
+
)
|
|
3497
|
+
arguments = tuple(sys.argv[1:] if argv is None else argv)
|
|
3498
|
+
try:
|
|
3499
|
+
dispatch = parse_session_dispatch(arguments)
|
|
3500
|
+
except SessionWorkerError as error:
|
|
3501
|
+
sys.stderr.write(f"{error.code}: {error.detail}\n")
|
|
3502
|
+
return 2
|
|
3503
|
+
try:
|
|
3504
|
+
config = config_from_environment()
|
|
3505
|
+
summary = run_session_job(config=config, dispatch=dispatch)
|
|
3506
|
+
except SessionWorkerError as error:
|
|
3507
|
+
if error.code in CONSUMED_SESSION_DISPATCH_CODES:
|
|
3508
|
+
LOGGER.info(
|
|
3509
|
+
"session dispatch %s was already consumed; exiting without work", error.code
|
|
3510
|
+
)
|
|
3511
|
+
return 0
|
|
3512
|
+
sys.stderr.write(f"{error.code}: {error.detail}\n")
|
|
3513
|
+
return 1
|
|
3514
|
+
except HostedWorkerError as error:
|
|
3515
|
+
sys.stderr.write(f"{error.code}: {error.detail}\n")
|
|
3516
|
+
return 1
|
|
3517
|
+
# A task that reaches Studio's terminal session state completed its work even when the state
|
|
3518
|
+
# is ``failed``: the dispatch was consumed and a retry must not create a second attempt.
|
|
3519
|
+
if summary["status"] == "stopped":
|
|
3520
|
+
return 0
|
|
3521
|
+
failure = summary.get("failure_code") or "SESSION_WORKER_FAILED"
|
|
3522
|
+
sys.stderr.write(f"{failure}: session job did not reach a terminal session state\n")
|
|
3523
|
+
return 1
|
|
3524
|
+
|
|
3525
|
+
|
|
3526
|
+
def session_worker_main(argv: Sequence[str] | None = None) -> int:
|
|
3527
|
+
"""The legacy Cloud Run service entrypoint, retained unchanged for rollback."""
|
|
3528
|
+
|
|
3529
|
+
logging.basicConfig(
|
|
3530
|
+
level=logging.INFO, stream=sys.stderr, format="%(levelname)s %(name)s %(message)s"
|
|
3531
|
+
)
|
|
3532
|
+
# This executable is also the fresh interpreter launched by
|
|
3533
|
+
# ``_LegacySessionProcess``. In that path Python invokes this module with
|
|
3534
|
+
# ``--legacy-bootstrap-fd <fd>`` in ``sys.argv``; treating a default ``None``
|
|
3535
|
+
# as an empty sequence would incorrectly enter ``serve`` a second time and
|
|
3536
|
+
# contend with the parent service for ``PORT``. Keep an explicitly supplied
|
|
3537
|
+
# empty sequence meaningful for unit callers, but mirror the Job entrypoint
|
|
3538
|
+
# when invoked as a console script.
|
|
3539
|
+
arguments = tuple(sys.argv[1:] if argv is None else argv)
|
|
3540
|
+
if arguments and arguments[0] == LEGACY_BOOTSTRAP_FD_ARGUMENT:
|
|
3541
|
+
return _legacy_bootstrap_child_main(arguments)
|
|
3542
|
+
if arguments:
|
|
3543
|
+
sys.stderr.write("mr-data-session-worker takes no arguments\n")
|
|
3544
|
+
return 2
|
|
3545
|
+
try:
|
|
3546
|
+
config = config_from_environment()
|
|
3547
|
+
except (SessionWorkerError, HostedWorkerError) as error:
|
|
3548
|
+
sys.stderr.write(f"{error.code}: {error.detail}\n")
|
|
3549
|
+
return 1
|
|
3550
|
+
return serve(config)
|
|
3551
|
+
|
|
3552
|
+
|
|
3553
|
+
if __name__ == "__main__": # pragma: no cover - exercised as a console script
|
|
3554
|
+
raise SystemExit(session_worker_main())
|