From 878c49f50fbb9820bce6b33a68bb5b8c785ce0ae Mon Sep 17 00:00:00 2001 From: jiangxt2 Date: Thu, 3 Sep 2026 23:42:24 +0800 Subject: [PATCH 1/2] feat(algorithms): unify PyTorch runtime contracts Signed-off-by: jiangxt2 --- docs/STABILITY.md | 5 +- docs/architecture/ray-first-torch-recipes.md | 64 +- docs/how-to/custom-distributed-algorithms.md | 85 +- docs/how-to/training.md | 9 +- docs/reference/api/algorithms-training.md | 225 +- docs/reference/api/data.md | 6 +- pyproject.toml | 2 +- src/tributo/_bootstrap.py | 6 +- src/tributo/algorithms/__init__.py | 130 +- src/tributo/algorithms/api/__init__.py | 90 + src/tributo/algorithms/api/distribution.py | 774 +++- src/tributo/algorithms/api/execution.py | 502 ++ src/tributo/algorithms/api/models.py | 42 +- src/tributo/algorithms/api/support.py | 36 + src/tributo/algorithms/api/torch_runtime.py | 2246 +++++++++ src/tributo/algorithms/composition.py | 8 +- src/tributo/algorithms/core/builder.py | 210 +- src/tributo/algorithms/core/dispatcher.py | 285 +- src/tributo/algorithms/core/planner.py | 25 +- src/tributo/algorithms/spi/__init__.py | 44 +- src/tributo/algorithms/spi/execution.py | 19 + src/tributo/algorithms/spi/torch.py | 705 ++- src/tributo/inference/kernel.py | 22 + .../algorithm_inputs/ingestion.py | 18 + .../algorithm_runtimes/collective.py | 22 +- .../algorithm_runtimes/portable_metrics.py | 11 +- .../algorithm_runtimes/ray_data_config.py | 16 +- .../algorithm_runtimes/ray_train_torch.py | 4028 +++++++++++++++++ .../algorithm_runtimes/torch_recipe.py | 1483 ------ src/tributo/integrations/sources/__init__.py | 4 +- src/tributo/integrations/sources/ray_torch.py | 542 +++ .../integrations/sources/ray_torch_recipe.py | 194 - src/tributo/plugin.py | 19 +- .../test_official_algorithm_migration.py | 2 +- tests/algorithms/test_portable_models.py | 4 +- .../algorithms/test_torch_recipe_contract.py | 43 +- tests/algorithms/test_torch_recipe_plugin.py | 4 +- tests/algorithms/test_torch_recipe_worker.py | 318 +- .../test_torch_runtime_contract_v1.py | 920 ++++ .../test_training_recipe_v2_contract.py | 98 +- .../__init__.py | 123 +- tests/inference/test_contracts.py | 50 + .../test_distributed_algorithm_it_contract.py | 4 +- tests/support/torch_recipe.py | 113 +- tests/support/training_recipe_v2.py | 155 +- tests/test_plugin.py | 6 +- tests/test_stability_inventory.py | 3 +- .../exporters/test_architecture_contracts.py | 2 +- .../exporters/test_torch_recipe_source.py | 292 +- .../jobs/official_algorithm_gate_job.py | 43 +- .../jobs/official_algorithm_matrix.py | 20 +- tests/training/test_dnn_pu_training.py | 4 +- tests/training/test_training_lifecycle.py | 2 +- tools/check_public_api_annotations.py | 6 + 54 files changed, 11566 insertions(+), 2523 deletions(-) create mode 100644 src/tributo/algorithms/api/torch_runtime.py create mode 100644 src/tributo/integrations/algorithm_runtimes/ray_train_torch.py delete mode 100644 src/tributo/integrations/algorithm_runtimes/torch_recipe.py create mode 100644 src/tributo/integrations/sources/ray_torch.py delete mode 100644 src/tributo/integrations/sources/ray_torch_recipe.py create mode 100644 tests/algorithms/test_torch_runtime_contract_v1.py diff --git a/docs/STABILITY.md b/docs/STABILITY.md index 88d01ef..fd103d4 100644 --- a/docs/STABILITY.md +++ b/docs/STABILITY.md @@ -97,6 +97,7 @@ from the legacy setup-only propagation rule. | `tributo.algorithms.api.context` — `UserExecutionContext` | `alpha` | Restricted context for trusted module-qualified Worker functions | | `tributo.algorithms.api.errors` | `alpha` | Portable execution error taxonomy | | `tributo.algorithms.api.support` | `alpha` | Trusted Wheel support evidence, execution semantics, expiry, and revocation | +| `tributo.algorithms.api.torch_runtime` | `alpha` | Versioned Torch Runtime helpers, Stage identity, checkpoint and reducer contracts | | `tributo.algorithms.conformance` | `alpha` | Descriptor-only and installed algorithm Wheel Conformance Testkit | | `tributo.algorithms.builtin.*` | `deprecated` | Production algorithms moved to the official `tributo-algorithms` Wheels; Core retains only public SPI and Ray runtimes | | `tributo.algorithms.core.builder` — `AlgorithmBuilder` | `alpha` | Provisional sklearn and Custom Ray Function registration builders | @@ -105,7 +106,7 @@ from the legacy setup-only propagation rule. | `tributo.algorithms.spi.execution` | `alpha` | Provisional operation and Runtime execution protocols | | `tributo.algorithms.spi.contracts` | `alpha` | Executable algorithm contract validator protocol | | `tributo.algorithms.spi.input` | `alpha` | Two-stage input resolution and Driver/Worker ownership contracts | -| `tributo.algorithms.spi.torch` | `alpha` | Narrow model/loss/optimizer/metric recipe contract lowered to Ray Train | +| `tributo.algorithms.spi.torch` | `alpha` | Versioned TorchRecipe and RayTorchAdapter contracts | ### Data (tributo.data.*) @@ -202,7 +203,7 @@ from the legacy setup-only propagation rule. | `tributo.integrations.validators.*` | `beta` | Built-in validator implementations | | `tributo.integrations.sources` | `beta` | Built-in source provider package | | `tributo.integrations.sources.*` | `beta` | Built-in source providers | -| `tributo.integrations.sources.ray_torch_recipe` | `alpha` | Generic trusted Torch recipe checkpoint provider | +| `tributo.integrations.sources.ray_torch` | `alpha` | Generic trusted Torch checkpoint provider | | `tributo.integrations.storage` | `beta` | Built-in storage adapter package | | `tributo.integrations.storage.*` | `beta` | Built-in storage backends | | `tributo.integrations.hooks` | `beta` | Built-in Hook package | diff --git a/docs/architecture/ray-first-torch-recipes.md b/docs/architecture/ray-first-torch-recipes.md index 44e4b9a..cdedf74 100644 --- a/docs/architecture/ray-first-torch-recipes.md +++ b/docs/architecture/ray-first-torch-recipes.md @@ -1,11 +1,10 @@ # Ray-first torch recipes -`TorchTrainingRecipe` is the low-code path for scalar-column dense tabular -PyTorch models. -An algorithm package defines model, loss, optimizer, and metric factories. -Tributo lowers that recipe to the existing formal collective runtime and keeps -deployment, data movement, distributed coordination, checkpoint upload, and -Bundle publication out of user code. +Tributo exposes one versioned `TorchRecipe` for Core-owned training loops and +one `RayTorchAdapter` for framework-owned loops. Both are selected by the +`RAY_TRAIN_TORCH` strategy and executed by `tributo.ray_train_torch`. +`TorchPolicy.execution_plan` is the single source of truth for routing, Stage +order, checkpoint dependencies, state layout and final export Stage. ## Ray reuse audit @@ -21,8 +20,8 @@ implementation and its public annotations, not unversioned latest docs. | Torch batches | Public `DataIterator.iter_torch_batches()` | Validate feature/label roles and pass bounded batch options | | Device and DDP | Stable `ray.train.torch.prepare_model()` | Reject unvalidated BatchNorm and GPU claims | | Metrics | PyTorch tensors and `torch.distributed` collectives | Apply the descriptor's existing `MetricReduction`; do not infer reducers from names | -| Checkpoint transfer | Ray `Checkpoint`, `train.report()`, `RunConfig`, and retention | Write bounded model/optimizer/RNG metadata and validate epoch-boundary resume identity | -| Run identity | V2 `RunConfig.name` and storage context | Derive a unique Ray run name from the existing Tributo `run_id` so independent runs cannot inherit each other's checkpoints | +| Checkpoint transfer | Ray `Checkpoint`, `train.report()`, `RunConfig`, and retention | Write bounded model, optimizer, scheduler, gradient scaling, and RNG metadata at complete optimizer-window boundaries and validate Stage identity | +| Run identity | V2 `RunConfig.name` and storage context | Derive a unique Ray run name from `run_id`, `invocation_id`, `stage_id`, and the Policy/execution-plan identity digest so retries reuse only the matching Stage directory | | Export and serving | Existing Torch ONNX exporter, ONNX Runtime validator, Bundle publisher, reader, batch inference, and Serve flavor | Reconstruct the trusted recipe model and create one `ExportSource` | | Local execution | `ray.init(address="local")` | Own and close only the runtime created by Tributo | | Existing or provisioned cluster | `ray.init(address="auto")`, Ray Jobs, KubeRay RayJob, or Cluster Launcher | Attach or expose the same entrypoint; workload code never provisions or deletes the cluster | @@ -34,12 +33,18 @@ changes only the `equal` argument while preserving Ray execution options, training-resource exclusion, and locality hints. A future Ray public option should replace this adapter. -V2 rejects a non-null `resume_from_checkpoint` argument as deprecated. Worker -recovery within one Ray run still uses `train.get_checkpoint()`. For an -explicit new-run resume, Tributo passes the validated, worker-visible epoch -checkpoint path to the recipe and restores model, optimizer, and RNG state in -the worker. This does not replace Ray's checkpoint upload, retention, or -failure-recovery controller. +Ray Train V2 rejects a non-null `resume_from_checkpoint` argument as deprecated. +The Core Runtime never supplies it: failure retry uses `train.get_checkpoint()`; +cross-Stage and cross-Run recovery use a credential-free Core Worker control +envelope. A full `TorchRecoveryEnvelope` records completed Stage locators and an +optional active Stage; completed Stage evidence is persisted in the validated +checkpoint sidecar so a recovery-only invocation can produce the same Receipt. +A typed progress manifest preserves per-rank coverage/metric prefixes, reducer +observation, Dataset cursor and whether the current epoch's Scheduler step has +already run. Remote Stage payloads remain unavailable under a deterministic +staging prefix until a matching commit marker is written. +A Torch-only preflight and invocation-local one-shot lease complete before Ray +or input resources are opened. ## Uneven input decision @@ -54,11 +59,19 @@ The first implementation therefore uses a narrow dynamic alignment protocol: - all ranks exchange only the active row count before each step; - exhausted ranks run a zero-row DDP forward/backward and contribute zero gradient while active ranks retain every observed row; -- loss scaling uses the global active row count, so uneven final batches are - sample weighted; +- loss scaling uses the algorithm-declared global normalizer for the complete + accumulation window, so uneven final batches remain mathematically weighted + by the declared unit (rows, sample weights, or valid tokens); - no row is dropped or replayed; -- all ranks report once per epoch and only rank zero attaches the replicated - checkpoint. +- all ranks report at the declared `checkpoint_interval_windows` cadence (one + complete optimizer window by default), and only the checkpoint-owner rank + attaches the replicated checkpoint. + +Optional validation and test roles use the same rank-alignment rule. When a +present role has fewer rows than Workers, exhausted ranks execute a typed +zero-row validation step and contribute zero metrics; all ranks still perform +the same metric collectives. A role that is absent remains explicitly absent, +while an explicitly empty role is rejected by the route contract. This protocol supplements domain correctness. It does not replace Ray Dataset sharding, PyTorch DDP, process-group setup, or Ray checkpoint transport. Empty @@ -77,10 +90,17 @@ original column names as separate ONNX inputs so Bundle inference does not lose the declared `InputBinding` contract. Vector, sequence, jagged, and graph inputs remain later profiles rather than implicit shape inference. -Advanced models may override the bounded `forward()` and `compute_loss()` -methods. Models that need custom training steps, multiple optimizers, framework -callbacks, graph sampling, sharded embeddings, or framework-owned checkpoints -use a framework adapter or the existing full worker-loop SPI. +Advanced models implement the typed `TorchRecipe` hooks. Models that need +custom training steps, framework callbacks, graph sampling, sharded embeddings, +or framework-owned checkpoints use `RayTorchAdapter`; the Adapter still cannot +create a nested Trainer or declare a second execution plan. + +An Adapter declares its `TorchArtifactPlan` through the Core Provider. The +Provider attaches that plan to the Adapter's `ExportSource` before invoking the +single Core Bundle exporter, and rejects conflicting plan metadata. Adapter +worker configuration receives algorithm-owned values plus credential-free input +binding metadata; Core Ray paths, recovery locators, and output publication +URIs are never accepted in that JSON payload. Capabilities not described by this contract are not implied by the recipe, dependency set, or framework installation. diff --git a/docs/how-to/custom-distributed-algorithms.md b/docs/how-to/custom-distributed-algorithms.md index 478c4e1..969f666 100644 --- a/docs/how-to/custom-distributed-algorithms.md +++ b/docs/how-to/custom-distributed-algorithms.md @@ -217,63 +217,74 @@ not restrict the algorithm family. ## Use a low-code PyTorch recipe -For scalar-column dense tabular models, subclass `TorchTrainingRecipe` and implement only -the four factories. The runtime stacks the feature columns declared by -`InputBinding` into one float32 tensor and supplies the declared label and -optional sample-weight roles. - -For multi-worker exact coverage, an exhausted rank can invoke `forward()` with -a zero-row tensor while another rank consumes its final batch. The model must -accept an empty leading batch dimension and return the corresponding empty -output shape; no observed row is replayed as padding. +For Core-owned training, subclass `TorchRecipe` and implement the typed +`build_modules`, `adapt_batch`, `training_step`, `validation_step`, +`configure_optimizers`, `metric_plan`, and `artifact_plan` hooks. Framework- +owned training instead subclasses `RayTorchAdapter`; the Adapter validates its +environment, binds role datasets, receives a Core-selected checkpoint context, +and cannot create a nested Trainer or declare a second execution plan. ```python -from collections.abc import Mapping -from typing import Any - -from tributo.algorithms import TorchTrainingRecipe +from tributo.algorithms import ( + TorchBatch, + TorchLossContribution, + TorchMetricPlan, + TorchOptimizationPlan, + TorchRecipe, + TorchStepResult, +) -class BinaryLinearRecipe(TorchTrainingRecipe): - def model_factory(self, config: Mapping[str, Any]) -> object: +class BinaryLinearRecipe(TorchRecipe): + def build_modules(self, context): import torch - return torch.nn.Linear(int(config["input_features"]), 1) + return {"model": torch.nn.Linear(2, 1), "loss": torch.nn.BCEWithLogitsLoss()} - def loss_factory(self, config: Mapping[str, Any]) -> object: + def adapt_batch(self, batch, context): import torch - return torch.nn.BCEWithLogitsLoss() + features = torch.as_tensor(batch["features"]) + targets = torch.as_tensor(batch["label"]) + return TorchBatch(positional=(features,), targets=targets, local_rows=len(targets)) - def optimizer_factory(self, model: object, config: Mapping[str, Any]) -> object: + def training_step(self, modules, batch, context): import torch - return torch.optim.Adam(model.parameters(), lr=float(config["lr"])) + predictions = modules["model"](batch.positional[0]) + numerator = torch.nn.functional.binary_cross_entropy_with_logits( + predictions, batch.targets.float(), reduction="sum" + ) + return TorchStepResult( + outputs={"prediction": predictions}, + loss=TorchLossContribution(numerator, batch.local_rows), + ) - def metric_factories(self, config: Mapping[str, Any]): + def validation_step(self, modules, batch, context): + return self.training_step(modules, batch, context) + + def configure_optimizers(self, modules, context): import torch - return { - "accuracy": lambda prediction, target: ( - (torch.sigmoid(prediction) >= 0.5) == target.bool() - ) - } + return TorchOptimizationPlan(torch.optim.Adam(modules["model"].parameters(), lr=1e-3)) + + def metric_plan(self, context): + return TorchMetricPlan({"train_loss": "sum_count"}) + + def artifact_plan(self, context): + return {"source_kind": "torch_module", "roles": {"inference": "onnx-model"}} ``` -Declare the reducer and use `AlgorithmBuilder.from_torch_recipe()` to lower the -class to the existing Ray Train collective runtime. `train_loss` is added with -the fixed `sum_count` reducer. The supported reducer set is `sum_count`, -`weighted_mean`, `min`, and `max`; a weighted metric requires -`InputBinding.sample_weight_name`. +Use `AlgorithmBuilder.from_torch()` to lower a Recipe, or +`AlgorithmBuilder.from_torch_adapter()` for a framework Adapter. Losses always +submit explicit numerator/normalizer pairs; model-specific composite reducers +remain owned by the algorithm Wheel while Core owns collectives and scaling. The complete independent package used by conformance lives under -`tests/fixtures/torch_recipe_algorithm_plugin`. It contains no Tributo builtin, +the installed algorithm Wheel. It contains no Tributo private Runtime import, Ray worker loop, checkpoint upload, Bundle publisher, or deployment code. -Default recipes receive a pre-bound training Dataset and do not own train/test -splitting. Use the advanced recipe hooks only for a custom model invocation or -loss call. Choose a framework adapter or `CollectiveAlgorithm` when the model -needs a custom training step, multiple optimizers, framework callbacks, or -special collectives. +Default Recipes receive role-routed datasets; the Core Runtime owns sharding, +checkpoint handling, evidence and Bundle publication. ## Choose the state coordination strategy diff --git a/docs/how-to/training.md b/docs/how-to/training.md index 73b68a1..0a7cbe4 100644 --- a/docs/how-to/training.md +++ b/docs/how-to/training.md @@ -1,9 +1,9 @@ # Training on Ray Run XGBoost, X-Learner, DNN, and PU through real distributed state coordination. -XGBoost and X-Learner use framework-native coordination; formal DNN lowers a first-party -`TorchTrainingRecipe` to the common Ray-owned worker loop, while PU retains its -specialized Ray Train/PyTorch DDP kernel. A one-worker local run is supported +XGBoost and X-Learner use framework-native coordination; PyTorch algorithms use +the unified `TorchRecipe` or `RayTorchAdapter` contract and one Core-owned Ray +Train Torch Runtime. A one-worker local run is supported but is not reported as distributed training. Formal algorithms choose an explicit `local` or `cluster` execution profile. @@ -18,7 +18,8 @@ the Ray Jobs API on an isolated Docker cluster with two worker nodes. It proves deployment-independent sharding, state coordination, receipt, and Bundle semantics without making Docker, Kubernetes, or VM providers execution profiles. -Ordinary third-party PyTorch models should implement `TorchTrainingRecipe`. +Ordinary third-party PyTorch models should implement `TorchRecipe`; framework- +owned loops should implement `RayTorchAdapter`. Advanced packages may implement `CollectiveAlgorithm`, `MapReduceAlgorithm`, or `FrameworkNativeAlgorithm` and publish a `DistributedAlgorithmDescriptor` through the `tributo.algorithms` entry-point diff --git a/docs/reference/api/algorithms-training.md b/docs/reference/api/algorithms-training.md index 13519f0..8779c2a 100644 --- a/docs/reference/api/algorithms-training.md +++ b/docs/reference/api/algorithms-training.md @@ -56,6 +56,10 @@ documentation for every public stability tier. :no-members: ``` +```{autoclass} tributo.algorithms.api.distribution.ComponentStageTorchPlan +:no-members: +``` + ```{autoclass} tributo.algorithms.api.distribution.DistributedExactness :no-members: ``` @@ -104,6 +108,10 @@ documentation for every public stability tier. :no-members: ``` +```{autoclass} tributo.algorithms.api.distribution.SingleStageTorchPlan +:no-members: +``` + ```{autoclass} tributo.algorithms.api.distribution.StateCoordination :no-members: ``` @@ -112,6 +120,22 @@ documentation for every public stability tier. :no-members: ``` +```{autoclass} tributo.algorithms.api.distribution.TorchDatasetRoute +:no-members: +``` + +```{autoclass} tributo.algorithms.api.distribution.TorchExecutionPlan +:no-members: +``` + +```{autoclass} tributo.algorithms.api.distribution.TorchPolicy +:no-members: +``` + +```{autoclass} tributo.algorithms.api.distribution.TorchStageSpec +:no-members: +``` + ```{autoclass} tributo.algorithms.api.distribution.WorkerRange :no-members: ``` @@ -141,6 +165,10 @@ documentation for every public stability tier. ## `tributo.algorithms.api.execution` +```{autoclass} tributo.algorithms.api.execution.ComponentStageEvidence +:no-members: +``` + ```{autoclass} tributo.algorithms.api.execution.ExecutionReceipt :no-members: ``` @@ -149,10 +177,22 @@ documentation for every public stability tier. :no-members: ``` +```{autoclass} tributo.algorithms.api.execution.ReplicatedTorchStateEvidence +:no-members: +``` + ```{autoclass} tributo.algorithms.api.execution.StateCoordinationEvidence :no-members: ``` +```{autoclass} tributo.algorithms.api.execution.TorchExecutionEvidence +:no-members: +``` + +```{autoclass} tributo.algorithms.api.execution.TorchRoleExecutionEvidence +:no-members: +``` + ```{autoclass} tributo.algorithms.api.execution.WorkerExecutionEvidence :no-members: ``` @@ -276,6 +316,133 @@ documentation for every public stability tier. ``` +## `tributo.algorithms.api.torch_runtime` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchAccumulationWindow +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchBackwardContext +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchBackwardResult +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchCheckpointDescriptor +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchCheckpointLocator +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchCheckpointPayloadDraft +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchCheckpointProgress +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchCheckpointRef +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchCompositeGlobalState +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchCompositeLossContribution +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchGlobalLossContext +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchGlobalLossReducer +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchGlobalLossReduction +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchLossContribution +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchMetricContribution +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchMetricPolicy +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchMetricReductionContext +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchMetricReductionResult +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchPreflightLease +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchPreflightTokenData +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchRankProgressStatistics +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchRecoveryEnvelope +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchRuntimeExecutionEnvelope +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchStageRunIdentity +:no-members: +``` + +```{autoclass} tributo.algorithms.api.torch_runtime.TorchWorkerControlEnvelope +:no-members: +``` + +```{autofunction} tributo.algorithms.api.torch_runtime.apply_torch_loss_backward +``` + +```{autofunction} tributo.algorithms.api.torch_runtime.claim_torch_run_directory +``` + +```{autofunction} tributo.algorithms.api.torch_runtime.describe_torch_checkpoint +``` + +```{autofunction} tributo.algorithms.api.torch_runtime.invoke_torch_global_loss_reducer +``` + +```{autofunction} tributo.algorithms.api.torch_runtime.reduce_torch_metrics +``` + +```{autofunction} tributo.algorithms.api.torch_runtime.report_torch_checkpoint +``` + +```{autofunction} tributo.algorithms.api.torch_runtime.torch_run_config_name +``` + +```{autofunction} tributo.algorithms.api.torch_runtime.validate_torch_retry_identity +``` + + ## `tributo.algorithms.composition` ```{autofunction} tributo.algorithms.composition.build_algorithm_dispatcher @@ -378,6 +545,10 @@ documentation for every public stability tier. :no-members: ``` +```{autoclass} tributo.algorithms.spi.execution.TorchRuntimePreflight +:no-members: +``` + ```{autoclass} tributo.algorithms.spi.execution.Transformable :no-members: ``` @@ -436,23 +607,67 @@ documentation for every public stability tier. ## `tributo.algorithms.spi.torch` -```{autoclass} tributo.algorithms.spi.torch.MetricPlan +```{autoclass} tributo.algorithms.spi.torch.RayTorchAdapter +:no-members: +``` + +```{autoclass} tributo.algorithms.spi.torch.TorchArtifactContext +:no-members: +``` + +```{autoclass} tributo.algorithms.spi.torch.TorchArtifactPlan +:no-members: +``` + +```{autoclass} tributo.algorithms.spi.torch.TorchBatch +:no-members: +``` + +```{autoclass} tributo.algorithms.spi.torch.TorchBatchContext +:no-members: +``` + +```{autoclass} tributo.algorithms.spi.torch.TorchBuildContext +:no-members: +``` + +```{autoclass} tributo.algorithms.spi.torch.TorchCheckpointContext +:no-members: +``` + +```{autoclass} tributo.algorithms.spi.torch.TorchMetricPlan +:no-members: +``` + +```{autoclass} tributo.algorithms.spi.torch.TorchModuleSet +:no-members: +``` + +```{autoclass} tributo.algorithms.spi.torch.TorchOptimizationPlan +:no-members: +``` + +```{autoclass} tributo.algorithms.spi.torch.TorchRecipe +:no-members: +``` + +```{autoclass} tributo.algorithms.spi.torch.TorchRuntimeContext :no-members: ``` -```{autoclass} tributo.algorithms.spi.torch.OptimizationPlan +```{autoclass} tributo.algorithms.spi.torch.TorchStageContext :no-members: ``` -```{autoclass} tributo.algorithms.spi.torch.TorchTrainingRecipe +```{autoclass} tributo.algorithms.spi.torch.TorchStepContext :no-members: ``` -```{autoclass} tributo.algorithms.spi.torch.TrainingRecipeV2 +```{autoclass} tributo.algorithms.spi.torch.TorchStepResult :no-members: ``` -```{autoclass} tributo.algorithms.spi.torch.TrainingStepResult +```{autoclass} tributo.algorithms.spi.torch.TorchWorkerCheckpointContext :no-members: ``` diff --git a/docs/reference/api/data.md b/docs/reference/api/data.md index cf9c599..f4a9571 100644 --- a/docs/reference/api/data.md +++ b/docs/reference/api/data.md @@ -350,13 +350,13 @@ documentation for every public stability tier. ``` -## `tributo.integrations.sources.ray_torch_recipe` +## `tributo.integrations.sources.ray_torch` -```{autoclass} tributo.integrations.sources.ray_torch_recipe.RayTorchRecipeSourceProvider +```{autoclass} tributo.integrations.sources.ray_torch.RayTorchSourceProvider :no-members: ``` -```{autoclass} tributo.integrations.sources.ray_torch_recipe.TorchRecipeSourceOptions +```{autoclass} tributo.integrations.sources.ray_torch.TorchSourceOptions :no-members: ``` diff --git a/pyproject.toml b/pyproject.toml index f5fcc20..a896ef9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -184,7 +184,7 @@ hf-onnx-v1 = "tributo.integrations.exporters.hf_onnx:HuggingFaceONNXExporter" onnx-quantizer-v1 = "tributo.integrations.exporters.onnx_quantizer:ONNXQuantizer" [project.entry-points."tributo.source_providers"] -ray-torch-recipe-v1 = "tributo.integrations.sources.ray_torch_recipe:RayTorchRecipeSourceProvider" +ray-torch-v1 = "tributo.integrations.sources.ray_torch:RayTorchSourceProvider" hf-v1 = "tributo.integrations.sources.huggingface:HuggingFaceSourceProvider" [project.entry-points."tributo.validators"] diff --git a/src/tributo/_bootstrap.py b/src/tributo/_bootstrap.py index 5ab8b9d..475b53e 100644 --- a/src/tributo/_bootstrap.py +++ b/src/tributo/_bootstrap.py @@ -43,11 +43,9 @@ def first_party_export_plugins() -> tuple[ def first_party_source_providers() -> tuple[type[ExportSourceProvider], ...]: """Return built-in checkpoint providers without entry-point metadata.""" - from tributo.integrations.sources.ray_torch_recipe import ( - RayTorchRecipeSourceProvider, - ) + from tributo.integrations.sources.ray_torch import RayTorchSourceProvider - return (RayTorchRecipeSourceProvider,) + return (RayTorchSourceProvider,) def first_party_model_flavors() -> tuple[type[BundleModelFlavor], ...]: diff --git a/src/tributo/algorithms/__init__.py b/src/tributo/algorithms/__init__.py index 1943878..6549028 100644 --- a/src/tributo/algorithms/__init__.py +++ b/src/tributo/algorithms/__init__.py @@ -16,6 +16,8 @@ ArtifactFile, BackendInputCompatibility, CollectivePolicy, + ComponentStageEvidence, + ComponentStageTorchPlan, ContractBinding, ContractBindingSet, DistributedAlgorithmDescriptor, @@ -41,17 +43,59 @@ MetricReduction, ParallelEnsemblePolicy, QualifiedReference, + ReplicatedTorchStateEvidence, ResultPolicy, RuntimeBinding, RuntimeTopology, + SingleStageTorchPlan, StateCoordination, StateCoordinationEvidence, StateField, SupportTier, + TorchAccumulationWindow, + TorchBackwardContext, + TorchBackwardResult, + TorchCheckpointDescriptor, + TorchCheckpointLocator, + TorchCheckpointPayloadDraft, + TorchCheckpointProgress, + TorchCheckpointRef, + TorchCompositeGlobalState, + TorchCompositeLossContribution, + TorchDatasetRoute, + TorchExecutionEvidence, + TorchExecutionPlan, + TorchGlobalLossContext, + TorchGlobalLossReducer, + TorchGlobalLossReduction, + TorchLossContribution, + TorchMetricContribution, + TorchMetricPolicy, + TorchMetricReductionContext, + TorchMetricReductionResult, + TorchPolicy, + TorchPreflightLease, + TorchPreflightTokenData, + TorchRankProgressStatistics, + TorchRecoveryEnvelope, + TorchRoleExecutionEvidence, + TorchRuntimeExecutionEnvelope, + TorchStageRunIdentity, + TorchStageSpec, + TorchStepLoss, + TorchWorkerControlEnvelope, UserExecutionContext, WorkerExecutionEvidence, WorkerRange, WorkerResources, + apply_torch_loss_backward, + claim_torch_run_directory, + describe_torch_checkpoint, + invoke_torch_global_loss_reducer, + reduce_torch_metrics, + report_torch_checkpoint, + torch_run_config_name, + validate_torch_retry_identity, ) from tributo.algorithms.composition import build_algorithm_dispatcher from tributo.algorithms.conformance import ( @@ -64,12 +108,23 @@ EnsembleUnitSpec, IterativeOptimizationAlgorithm, JoblibEstimatorRecipe, - MetricPlan, - OptimizationPlan, ParallelEnsembleAlgorithm, - TorchTrainingRecipe, - TrainingRecipeV2, - TrainingStepResult, + RayTorchAdapter, + TorchArtifactContext, + TorchArtifactPlan, + TorchBatch, + TorchBatchContext, + TorchBuildContext, + TorchCheckpointContext, + TorchMetricPlan, + TorchModuleSet, + TorchOptimizationPlan, + TorchRecipe, + TorchRuntimeContext, + TorchStageContext, + TorchStepContext, + TorchStepResult, + TorchWorkerCheckpointContext, ) __all__ = [ @@ -85,11 +140,13 @@ "AlgorithmRunResult", "AlgorithmSupportEvidence", "AlgorithmSupportEvidenceRegistry", + "ComponentStageEvidence", "ArtifactDraft", "ArtifactDistributionMode", "ArtifactFile", "BackendInputCompatibility", "CollectivePolicy", + "ComponentStageTorchPlan", "ContractBinding", "ContractBindingSet", "DistributedExactness", @@ -102,6 +159,7 @@ "ExecutionProfile", "ExecutionReceipt", "ExecutionRequest", + "ReplicatedTorchStateEvidence", "FrameworkNativePolicy", "EnsembleUnitSpec", "ImplementationDescriptor", @@ -116,25 +174,77 @@ "InputCoverageContract", "MapReducePolicy", "MetricReduction", - "MetricPlan", - "OptimizationPlan", "ParallelEnsembleAlgorithm", "ParallelEnsemblePolicy", "QualifiedReference", "ResultPolicy", "RuntimeBinding", "RuntimeTopology", + "SingleStageTorchPlan", "StateCoordination", "StateCoordinationEvidence", + "TorchExecutionEvidence", + "TorchRoleExecutionEvidence", "StateField", "SupportTier", - "TorchTrainingRecipe", - "TrainingRecipeV2", - "TrainingStepResult", "UserExecutionContext", "WorkerExecutionEvidence", "WorkerRange", "WorkerResources", + "TorchDatasetRoute", + "TorchExecutionPlan", + "TorchPolicy", + "TorchStageSpec", + "RayTorchAdapter", + "TorchArtifactContext", + "TorchArtifactPlan", + "TorchBatch", + "TorchBatchContext", + "TorchBuildContext", + "TorchCheckpointContext", + "TorchMetricPlan", + "TorchModuleSet", + "TorchOptimizationPlan", + "TorchRecipe", + "TorchRuntimeContext", + "TorchStageContext", + "TorchStepContext", + "TorchStepResult", + "TorchWorkerCheckpointContext", + "TorchAccumulationWindow", + "TorchBackwardContext", + "TorchBackwardResult", + "TorchCheckpointPayloadDraft", + "TorchCheckpointProgress", + "TorchCheckpointDescriptor", + "TorchCheckpointLocator", + "TorchCheckpointRef", + "TorchCompositeGlobalState", + "TorchCompositeLossContribution", + "TorchGlobalLossContext", + "TorchGlobalLossReducer", + "TorchGlobalLossReduction", + "TorchLossContribution", + "TorchMetricContribution", + "TorchMetricPolicy", + "TorchMetricReductionContext", + "TorchMetricReductionResult", + "TorchPreflightLease", + "TorchPreflightTokenData", + "TorchRecoveryEnvelope", + "TorchRankProgressStatistics", + "TorchRuntimeExecutionEnvelope", + "TorchStageRunIdentity", + "TorchStepLoss", + "TorchWorkerControlEnvelope", + "apply_torch_loss_backward", + "claim_torch_run_directory", + "describe_torch_checkpoint", + "invoke_torch_global_loss_reducer", + "reduce_torch_metrics", + "report_torch_checkpoint", + "torch_run_config_name", + "validate_torch_retry_identity", "build_algorithm_dispatcher", "validate_algorithm_descriptor_conformance", "validate_installed_algorithm_package", diff --git a/src/tributo/algorithms/api/__init__.py b/src/tributo/algorithms/api/__init__.py index 948e6fb..873a272 100644 --- a/src/tributo/algorithms/api/__init__.py +++ b/src/tributo/algorithms/api/__init__.py @@ -12,6 +12,7 @@ from tributo.algorithms.api.descriptor import DistributedAlgorithmDescriptor from tributo.algorithms.api.distribution import ( CollectivePolicy, + ComponentStageTorchPlan, DistributedExactness, DistributionSpec, DistributionStrategy, @@ -24,8 +25,13 @@ MetricReduction, ParallelEnsemblePolicy, ResultPolicy, + SingleStageTorchPlan, StateCoordination, StateField, + TorchDatasetRoute, + TorchExecutionPlan, + TorchPolicy, + TorchStageSpec, WorkerRange, WorkerResources, ) @@ -37,9 +43,13 @@ AlgorithmResolutionError, ) from tributo.algorithms.api.execution import ( + ComponentStageEvidence, ExecutionReceipt, ExecutionRequest, + ReplicatedTorchStateEvidence, StateCoordinationEvidence, + TorchExecutionEvidence, + TorchRoleExecutionEvidence, WorkerExecutionEvidence, ) from tributo.algorithms.api.models import ( @@ -75,6 +85,42 @@ DistributedSemantics, SupportTier, ) +from tributo.algorithms.api.torch_runtime import ( + TorchAccumulationWindow, + TorchBackwardContext, + TorchBackwardResult, + TorchCheckpointDescriptor, + TorchCheckpointLocator, + TorchCheckpointPayloadDraft, + TorchCheckpointProgress, + TorchCheckpointRef, + TorchCompositeGlobalState, + TorchCompositeLossContribution, + TorchGlobalLossContext, + TorchGlobalLossReducer, + TorchGlobalLossReduction, + TorchLossContribution, + TorchMetricContribution, + TorchMetricPolicy, + TorchMetricReductionContext, + TorchMetricReductionResult, + TorchPreflightLease, + TorchPreflightTokenData, + TorchRankProgressStatistics, + TorchRecoveryEnvelope, + TorchRuntimeExecutionEnvelope, + TorchStageRunIdentity, + TorchStepLoss, + TorchWorkerControlEnvelope, + apply_torch_loss_backward, + claim_torch_run_directory, + describe_torch_checkpoint, + invoke_torch_global_loss_reducer, + reduce_torch_metrics, + report_torch_checkpoint, + torch_run_config_name, + validate_torch_retry_identity, +) __all__ = [ "AlgorithmConfigurationError", @@ -98,6 +144,7 @@ "ArtifactFile", "BackendInputCompatibility", "CollectivePolicy", + "ComponentStageTorchPlan", "ContractBinding", "ContractBindingSet", "DistributedExactness", @@ -110,6 +157,10 @@ "ExecutionProfile", "ExecutionReceipt", "ExecutionRequest", + "ComponentStageEvidence", + "ReplicatedTorchStateEvidence", + "TorchExecutionEvidence", + "TorchRoleExecutionEvidence", "FailureCategory", "FrameworkNativePolicy", "ImplementationDescriptor", @@ -133,11 +184,50 @@ "StateCoordination", "StateCoordinationEvidence", "StateField", + "SingleStageTorchPlan", "SupportTier", "UserExecutionContext", "WorkerExecutionEvidence", "WorkerExecutionResult", "WorkerRange", "WorkerResources", + "TorchDatasetRoute", + "TorchExecutionPlan", + "TorchPolicy", + "TorchStageSpec", "canonical_digest", + "TorchAccumulationWindow", + "TorchBackwardContext", + "TorchBackwardResult", + "TorchCheckpointPayloadDraft", + "TorchCheckpointProgress", + "TorchCheckpointDescriptor", + "TorchCheckpointLocator", + "TorchCheckpointRef", + "TorchCompositeGlobalState", + "TorchCompositeLossContribution", + "TorchGlobalLossContext", + "TorchGlobalLossReducer", + "TorchGlobalLossReduction", + "TorchLossContribution", + "TorchMetricContribution", + "TorchMetricPolicy", + "TorchMetricReductionContext", + "TorchMetricReductionResult", + "TorchPreflightLease", + "TorchPreflightTokenData", + "TorchRecoveryEnvelope", + "TorchRankProgressStatistics", + "TorchRuntimeExecutionEnvelope", + "TorchStageRunIdentity", + "TorchStepLoss", + "TorchWorkerControlEnvelope", + "apply_torch_loss_backward", + "claim_torch_run_directory", + "describe_torch_checkpoint", + "invoke_torch_global_loss_reducer", + "reduce_torch_metrics", + "report_torch_checkpoint", + "torch_run_config_name", + "validate_torch_retry_identity", ] diff --git a/src/tributo/algorithms/api/distribution.py b/src/tributo/algorithms/api/distribution.py index d3247a5..5e81792 100644 --- a/src/tributo/algorithms/api/distribution.py +++ b/src/tributo/algorithms/api/distribution.py @@ -15,7 +15,7 @@ from collections.abc import Mapping from dataclasses import dataclass, field from enum import Enum -from typing import Any +from typing import Any, cast from tributo._common.immutable import FrozenDict from tributo.algorithms.api.errors import AlgorithmConfigurationError @@ -25,6 +25,7 @@ r"^[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*" r":[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*$" ) +_NAMESPACED_ID = re.compile(r"^[a-z][a-z0-9_.-]*$") def _positive_integer(value: object, field_name: str) -> int: @@ -53,6 +54,13 @@ def _string(value: object, field_name: str) -> str: return value +def _require_namespaced_id(value: object, field_name: str) -> None: + if not isinstance(value, str) or _NAMESPACED_ID.fullmatch(value) is None: + raise AlgorithmConfigurationError( + f"{field_name} must be a lower-case namespaced identifier" + ) + + def _mapping(value: object, field_name: str) -> Mapping[Any, Any]: if not isinstance(value, Mapping): raise AlgorithmConfigurationError(f"{field_name} must be a mapping") @@ -121,7 +129,7 @@ class DistributionStrategy(str, Enum): RAY_JOBLIB_ESTIMATOR = "ray_joblib_estimator" RAY_PARALLEL_ENSEMBLE = "ray_parallel_ensemble" RAY_ITERATIVE_OPTIMIZATION = "ray_iterative_optimization" - RAY_TRAIN_RECIPE_V2 = "ray_train_recipe_v2" + RAY_TRAIN_TORCH = "ray_train_torch" @PublicAPI(stability="alpha") @@ -129,6 +137,7 @@ class InputDistribution(str, Enum): """How training input reaches workers.""" SHARDED = "sharded" + ROLE_ROUTED = "role_routed" FRAMEWORK_OWNED = "framework_owned" FULL_DATASET = "full_dataset" @@ -138,6 +147,7 @@ class StateCoordination(str, Enum): """How worker-local state becomes one global model.""" ALL_REDUCE = "all_reduce" + TORCH_MANAGED = "torch_managed" FRAMEWORK_NATIVE = "framework_native" ASSOCIATIVE_REDUCE = "associative_reduce" ESTIMATOR_INTERNAL = "estimator_internal" @@ -510,6 +520,569 @@ def __post_init__(self) -> None: ) +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchDatasetRoute: + """Role-specific bounded input routing for one Torch Policy.""" + + role: str + mode: str + required: bool = True + min_total_rows_if_present: int = 1 + min_rows_per_worker: int = 1 + empty_rank_policy: str = "reject" + max_rows: int | None = None + max_bytes_per_worker: int | None = None + + def __post_init__(self) -> None: + _string(self.role, "Torch route role") + if self.mode not in {"split_exact", "replicate", "split_framework"}: + raise AlgorithmConfigurationError("invalid Torch route mode") + _boolean(self.required, "Torch route required") + _non_negative_integer( + self.min_total_rows_if_present, "Torch route minimum total rows" + ) + _non_negative_integer( + self.min_rows_per_worker, "Torch route minimum rows per worker" + ) + if self.empty_rank_policy not in {"reject", "zero_contribution"}: + raise AlgorithmConfigurationError("invalid Torch empty rank policy") + if ( + self.mode == "split_exact" + and self.required + and ( + self.min_total_rows_if_present < 1 + or self.min_rows_per_worker < 1 + or self.empty_rank_policy != "reject" + ) + ): + raise AlgorithmConfigurationError( + "required split_exact training routes must reject empty ranks" + ) + if ( + not self.required + and self.mode == "split_exact" + and ( + self.min_total_rows_if_present < 1 + or self.min_rows_per_worker != 0 + or self.empty_rank_policy != "zero_contribution" + ) + ): + raise AlgorithmConfigurationError( + "optional split_exact routes must use zero-contribution empty ranks" + ) + if self.empty_rank_policy == "zero_contribution" and self.required: + raise AlgorithmConfigurationError( + "zero-contribution routes must be optional evaluation roles" + ) + if self.mode == "replicate": + _positive_integer(self.max_rows, "Torch replicate max_rows") + _positive_integer( + self.max_bytes_per_worker, + "Torch replicate max_bytes_per_worker", + ) + elif self.max_rows is not None or self.max_bytes_per_worker is not None: + raise AlgorithmConfigurationError( + "replication budgets are only valid for replicate routes" + ) + + def to_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "role": self.role, + "mode": self.mode, + "required": self.required, + "min_total_rows_if_present": self.min_total_rows_if_present, + "min_rows_per_worker": self.min_rows_per_worker, + "empty_rank_policy": self.empty_rank_policy, + "max_rows": self.max_rows, + "max_bytes_per_worker": self.max_bytes_per_worker, + } + return payload + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchStageSpec: + """One Core-orchestrated stage in a Torch execution plan.""" + + stage_id: str + worker_loop_ref: str + input_roles: tuple[str, ...] + depends_on: tuple[str, ...] = () + checkpoint_from_stage: str | None = None + metric_mapping: Mapping[str, str] = field(default_factory=dict) + checkpoint_required: bool = True + checkpoint_interval_windows: int = 1 + + def __post_init__(self) -> None: + _string(self.stage_id, "Torch stage_id") + _qualified_reference(self.worker_loop_ref, "Torch worker_loop_ref") + roles = tuple(self.input_roles) + if not roles or any(not isinstance(role, str) or not role for role in roles): + raise AlgorithmConfigurationError("Torch stage input_roles are required") + if len(set(roles)) != len(roles): + raise AlgorithmConfigurationError("Torch stage input_roles must be unique") + depends = tuple(self.depends_on) + if any(not isinstance(item, str) or not item for item in depends): + raise AlgorithmConfigurationError("Torch stage dependencies are invalid") + if len(set(depends)) != len(depends): + raise AlgorithmConfigurationError("Torch stage dependencies must be unique") + if self.checkpoint_from_stage is not None and not isinstance( + self.checkpoint_from_stage, str + ): + raise AlgorithmConfigurationError("Torch checkpoint_from_stage is invalid") + _boolean(self.checkpoint_required, "Torch checkpoint_required") + _positive_integer( + self.checkpoint_interval_windows, + "Torch checkpoint_interval_windows", + ) + if any( + not isinstance(name, str) + or not name + or not isinstance(value, str) + or not value + for name, value in self.metric_mapping.items() + ): + raise AlgorithmConfigurationError( + "Torch stage metric_mapping must be named strings" + ) + if len(set(self.metric_mapping.values())) != len(self.metric_mapping): + raise AlgorithmConfigurationError( + "Torch stage metric_mapping targets must be unique" + ) + object.__setattr__(self, "input_roles", roles) + object.__setattr__(self, "depends_on", depends) + object.__setattr__( + self, "metric_mapping", FrozenDict(dict(self.metric_mapping)) + ) + + def to_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "stage_id": self.stage_id, + "worker_loop_ref": self.worker_loop_ref, + "input_roles": list(self.input_roles), + "depends_on": list(self.depends_on), + "checkpoint_from_stage": self.checkpoint_from_stage, + "metric_mapping": dict(self.metric_mapping), + "checkpoint_required": self.checkpoint_required, + "checkpoint_interval_windows": self.checkpoint_interval_windows, + } + return payload + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchExecutionPlan: + """Versioned closed union base for single and component Torch plans.""" + + api_version: int = 1 + + def to_dict(self) -> dict[str, Any]: + return {"api_version": self.api_version} + + @property + def stages(self) -> tuple[TorchStageSpec, ...]: + stage = getattr(self, "stage", None) + return (stage,) if isinstance(stage, TorchStageSpec) else () + + @property + def final_stage_id(self) -> str: + explicit = getattr(self, "_final_stage_id", None) + if isinstance(explicit, str) and explicit: + return explicit + stages = self.stages + return stages[-1].stage_id if len(stages) == 1 else "" + + @property + def digest(self) -> str: + return hashlib.sha256( + json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class SingleStageTorchPlan(TorchExecutionPlan): + """One-stage Torch execution plan.""" + + stage: TorchStageSpec = field( + default_factory=lambda: TorchStageSpec( + "train", + "tributo.integrations.algorithm_runtimes.ray_train_torch:" + "torch_recipe_train_loop_per_worker", + ("train",), + ) + ) + + def __post_init__(self) -> None: + if self.api_version != 1: + raise AlgorithmConfigurationError( + "Torch execution plan api_version must be 1" + ) + if self.stage.depends_on or self.stage.checkpoint_from_stage is not None: + raise AlgorithmConfigurationError( + "single Torch stage cannot have dependencies" + ) + + def to_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "kind": "single", + "api_version": self.api_version, + "stages": [self.stage.to_dict()], + "final_stage_id": self.stage.stage_id, + } + return payload + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class ComponentStageTorchPlan(TorchExecutionPlan): + """Ordered multi-stage Torch execution plan.""" + + stages: tuple[TorchStageSpec, ...] = () + final_stage_id: str = "" + + def __post_init__(self) -> None: + if self.api_version != 1 or not self.stages: + raise AlgorithmConfigurationError("component Torch plan requires stages") + stages = tuple(self.stages) + ids = [stage.stage_id for stage in stages] + if len(set(ids)) != len(ids) or self.final_stage_id not in ids: + raise AlgorithmConfigurationError("component Torch stage IDs are invalid") + prior: set[str] = set() + for stage in stages: + if any(dep not in prior for dep in stage.depends_on): + raise AlgorithmConfigurationError( + "Torch stage dependencies must reference earlier stages" + ) + if ( + stage.checkpoint_from_stage is not None + and stage.checkpoint_from_stage not in prior + ): + raise AlgorithmConfigurationError( + "Torch checkpoint_from_stage must reference an earlier stage" + ) + if stage.checkpoint_from_stage is not None: + source = next( + item + for item in stages + if item.stage_id == stage.checkpoint_from_stage + ) + if not source.checkpoint_required: + raise AlgorithmConfigurationError( + "Torch checkpoint_from_stage requires a checkpoint-producing source" + ) + prior.add(stage.stage_id) + object.__setattr__(self, "stages", stages) + + def to_dict(self) -> dict[str, Any]: + return { + "kind": "component", + "api_version": self.api_version, + "stages": [stage.to_dict() for stage in self.stages], + "final_stage_id": self.final_stage_id, + } + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchPolicy: + """Versioned policy carrying all Torch routing and execution facts.""" + + torch_runtime_api_version: int + loop_owner: str + parallelism_id: str + dataset_routing: tuple[TorchDatasetRoute, ...] + execution_plan: TorchExecutionPlan + state_layout: str + metric_reducers: Mapping[str, MetricReduction] + backend: str = "auto" + checkpoint_owner_rank: int = 0 + resume_supported: bool = True + same_world_size_resume: bool | None = True + rank_seeded: bool = True + checkpoint_adapter_ref: str | None = None + evidence_adapter_ref: str | None = None + global_loss_reducer_ref: str | None = None + global_loss_reducer_api_version: int | None = None + global_loss_reducer_code_digest: str | None = None + composite_loss_schema_id: str | None = None + capabilities: tuple[str, ...] = () + max_replicated_bytes_per_worker: int | None = None + + def __post_init__(self) -> None: + if self.torch_runtime_api_version != 1: + raise AlgorithmConfigurationError("Torch Runtime API version must be 1") + if self.loop_owner not in {"core_recipe", "adapter"}: + raise AlgorithmConfigurationError("Torch loop_owner is invalid") + _require_namespaced_id(self.parallelism_id, "Torch parallelism_id") + _qualified_reference( + self.checkpoint_adapter_ref, "checkpoint_adapter_ref" + ) if self.checkpoint_adapter_ref else None + _qualified_reference( + self.evidence_adapter_ref, "evidence_adapter_ref" + ) if self.evidence_adapter_ref else None + routes = tuple(self.dataset_routing) + if not routes or any( + not isinstance(route, TorchDatasetRoute) for route in routes + ): + raise AlgorithmConfigurationError("Torch Policy requires dataset routes") + if len({route.role for route in routes}) != len(routes): + raise AlgorithmConfigurationError("Torch Policy route roles must be unique") + if not isinstance(self.execution_plan, TorchExecutionPlan): + raise AlgorithmConfigurationError("Torch Policy requires an execution plan") + declared_stages = ( + (self.execution_plan.stage,) + if isinstance(self.execution_plan, SingleStageTorchPlan) + else self.execution_plan.stages + ) + stage_roles = {role for stage in declared_stages for role in stage.input_roles} + route_roles = {route.role for route in routes} + missing_roles = sorted(stage_roles - route_roles) + if missing_roles: + raise AlgorithmConfigurationError( + f"Torch Policy is missing routes for stage role(s): {missing_roles}" + ) + if route_roles - stage_roles: + raise AlgorithmConfigurationError( + "Torch Policy declares an input route unused by its execution plan" + ) + if any( + route.mode == "split_framework" and self.loop_owner != "adapter" + for route in routes + ): + raise AlgorithmConfigurationError( + "split_framework routing is only valid for Adapter-owned loops" + ) + replicated_budget = sum( + route.max_bytes_per_worker or 0 + for route in routes + if route.mode == "replicate" + ) + if replicated_budget and ( + self.max_replicated_bytes_per_worker is None + or replicated_budget > self.max_replicated_bytes_per_worker + ): + raise AlgorithmConfigurationError( + "Torch replicate routes exceed max_replicated_bytes_per_worker" + ) + if self.state_layout not in {"replicated", "component", "sharded"}: + raise AlgorithmConfigurationError("Torch Policy state_layout is invalid") + if self.state_layout == "replicated" and not isinstance( + self.execution_plan, SingleStageTorchPlan + ): + raise AlgorithmConfigurationError( + "replicated Torch state requires a single stage plan" + ) + if self.state_layout == "component" and not isinstance( + self.execution_plan, ComponentStageTorchPlan + ): + raise AlgorithmConfigurationError( + "component Torch state requires a component plan" + ) + if self.backend not in {"auto", "gloo", "nccl"}: + raise AlgorithmConfigurationError("Torch backend is invalid") + _non_negative_integer(self.checkpoint_owner_rank, "Torch checkpoint_owner_rank") + _boolean(self.resume_supported, "Torch resume_supported") + _boolean(self.rank_seeded, "Torch rank_seeded") + if self.same_world_size_resume is not None: + _boolean(self.same_world_size_resume, "Torch same_world_size_resume") + if self.resume_supported and self.same_world_size_resume is not True: + raise AlgorithmConfigurationError( + "Torch v1 recovery only supports same-world-size resume" + ) + if not self.resume_supported and self.same_world_size_resume is not None: + raise AlgorithmConfigurationError( + "unsupported Torch resume must omit same_world_size_resume" + ) + if self.max_replicated_bytes_per_worker is not None: + _positive_integer( + self.max_replicated_bytes_per_worker, + "Torch max_replicated_bytes_per_worker", + ) + reducer_fields = ( + self.global_loss_reducer_ref, + self.global_loss_reducer_api_version, + self.global_loss_reducer_code_digest, + self.composite_loss_schema_id, + ) + if any(value is not None for value in reducer_fields): + if not all(value is not None for value in reducer_fields): + raise AlgorithmConfigurationError( + "Torch composite reducer fields must be declared together" + ) + _qualified_reference( + cast(str, self.global_loss_reducer_ref), "global_loss_reducer_ref" + ) + if self.global_loss_reducer_api_version != 1: + raise AlgorithmConfigurationError( + "Torch global reducer API version must be 1" + ) + _digest_value = self.global_loss_reducer_code_digest + if ( + not isinstance(_digest_value, str) + or len(_digest_value) != 64 + or any(char not in "0123456789abcdef" for char in _digest_value) + ): + raise AlgorithmConfigurationError( + "Torch reducer code digest is invalid" + ) + if len(set(self.capabilities)) != len(self.capabilities): + raise AlgorithmConfigurationError("Torch capabilities must be unique") + for capability in self.capabilities: + _require_namespaced_id(capability, "Torch capability") + try: + metric_reducers = { + _string(name, "Torch metric name"): MetricReduction(value) + for name, value in self.metric_reducers.items() + } + except (TypeError, ValueError) as exc: + raise AlgorithmConfigurationError( + "Torch metric reducer is invalid" + ) from exc + if ( + "train_loss" not in metric_reducers + or metric_reducers["train_loss"] is not MetricReduction.SUM_COUNT + ): + raise AlgorithmConfigurationError( + "Torch metric_reducers must declare train_loss=sum_count" + ) + object.__setattr__(self, "dataset_routing", routes) + object.__setattr__(self, "metric_reducers", FrozenDict(metric_reducers)) + object.__setattr__(self, "capabilities", tuple(sorted(set(self.capabilities)))) + + def to_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "torch_runtime_api_version": self.torch_runtime_api_version, + "loop_owner": self.loop_owner, + "parallelism_id": self.parallelism_id, + "dataset_routing": [route.to_dict() for route in self.dataset_routing], + "execution_plan": self.execution_plan.to_dict(), + "state_layout": self.state_layout, + "metric_reducers": { + name: value.value + for name, value in sorted(self.metric_reducers.items()) + }, + "backend": self.backend, + "checkpoint_owner_rank": self.checkpoint_owner_rank, + "resume_supported": self.resume_supported, + "rank_seeded": self.rank_seeded, + "checkpoint_adapter_ref": self.checkpoint_adapter_ref, + "evidence_adapter_ref": self.evidence_adapter_ref, + "global_loss_reducer_ref": self.global_loss_reducer_ref, + "global_loss_reducer_api_version": self.global_loss_reducer_api_version, + "global_loss_reducer_code_digest": self.global_loss_reducer_code_digest, + "composite_loss_schema_id": self.composite_loss_schema_id, + "capabilities": list(self.capabilities), + "max_replicated_bytes_per_worker": self.max_replicated_bytes_per_worker, + } + if self.same_world_size_resume is not None: + payload["same_world_size_resume"] = self.same_world_size_resume + return payload + + @property + def digest(self) -> str: + payload = json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "TorchPolicy": + """Reconstruct a policy without importing an implementation module.""" + try: + plan_value = _mapping(value["execution_plan"], "Torch execution_plan") + raw_stages = _sequence(plan_value["stages"], "Torch execution stages") + stages = tuple( + TorchStageSpec( + stage_id=_string( + _mapping(item, "Torch stage")["stage_id"], "stage_id" + ), + worker_loop_ref=_string( + _mapping(item, "Torch stage")["worker_loop_ref"], + "worker_loop_ref", + ), + input_roles=tuple( + _sequence( + _mapping(item, "Torch stage")["input_roles"], "input_roles" + ) + ), + depends_on=tuple( + _sequence( + _mapping(item, "Torch stage").get("depends_on", ()), + "depends_on", + ) + ), + checkpoint_from_stage=_mapping(item, "Torch stage").get( + "checkpoint_from_stage" + ), + metric_mapping=_mapping( + _mapping(item, "Torch stage").get("metric_mapping", {}), + "metric_mapping", + ), + checkpoint_required=_boolean( + _mapping(item, "Torch stage").get("checkpoint_required", True), + "checkpoint_required", + ), + ) + for item in raw_stages + ) + if plan_value.get("kind") == "single": + execution_plan: TorchExecutionPlan = SingleStageTorchPlan( + api_version=plan_value["api_version"], stage=stages[0] + ) + elif plan_value.get("kind") == "component": + execution_plan = ComponentStageTorchPlan( + api_version=plan_value["api_version"], + stages=stages, + final_stage_id=plan_value["final_stage_id"], + ) + else: + raise AlgorithmConfigurationError( + "Torch execution plan kind is invalid" + ) + routes = tuple( + TorchDatasetRoute(**dict(_mapping(item, "Torch route"))) + for item in _sequence(value["dataset_routing"], "dataset_routing") + ) + metric_reducers = { + name: MetricReduction(reducer) + for name, reducer in _mapping( + value["metric_reducers"], "metric_reducers" + ).items() + } + return cls( + torch_runtime_api_version=value["torch_runtime_api_version"], + loop_owner=value["loop_owner"], + parallelism_id=value["parallelism_id"], + dataset_routing=routes, + execution_plan=execution_plan, + state_layout=value["state_layout"], + metric_reducers=metric_reducers, + backend=value.get("backend", "auto"), + checkpoint_owner_rank=value.get("checkpoint_owner_rank", 0), + resume_supported=value.get("resume_supported", True), + same_world_size_resume=value.get("same_world_size_resume"), + rank_seeded=value.get("rank_seeded", True), + checkpoint_adapter_ref=value.get("checkpoint_adapter_ref"), + evidence_adapter_ref=value.get("evidence_adapter_ref"), + global_loss_reducer_ref=value.get("global_loss_reducer_ref"), + global_loss_reducer_api_version=value.get( + "global_loss_reducer_api_version" + ), + global_loss_reducer_code_digest=value.get( + "global_loss_reducer_code_digest" + ), + composite_loss_schema_id=value.get("composite_loss_schema_id"), + capabilities=tuple(value.get("capabilities", ())), + max_replicated_bytes_per_worker=value.get( + "max_replicated_bytes_per_worker" + ), + ) + except (KeyError, TypeError, ValueError) as exc: + if isinstance(exc, AlgorithmConfigurationError): + raise + raise AlgorithmConfigurationError("invalid TorchPolicy payload") from exc + + StrategyPolicy = ( CollectivePolicy | MapReducePolicy @@ -517,6 +1090,7 @@ def __post_init__(self) -> None: | JoblibEstimatorPolicy | ParallelEnsemblePolicy | IterativeOptimizationPolicy + | TorchPolicy ) @@ -622,10 +1196,10 @@ def __post_init__(self) -> None: InputDistribution.SHARDED, StateCoordination.ITERATIVE_GLOBAL, ), - DistributionStrategy.RAY_TRAIN_RECIPE_V2: ( - CollectivePolicy, - InputDistribution.SHARDED, - StateCoordination.ALL_REDUCE, + DistributionStrategy.RAY_TRAIN_TORCH: ( + TorchPolicy, + InputDistribution.ROLE_ROUTED, + StateCoordination.TORCH_MANAGED, ), } policy_type, expected_input, expected_state = expected[strategy] @@ -650,6 +1224,15 @@ def __post_init__(self) -> None: raise AlgorithmConfigurationError( "checkpoint_owner_rank must fit every supported worker group" ) + if isinstance(self.policy, TorchPolicy): + maximum_rank = self.supported_worker_range.maximum + if ( + maximum_rank is not None + and self.policy.checkpoint_owner_rank >= maximum_rank + ): + raise AlgorithmConfigurationError( + "Torch checkpoint_owner_rank must fit every supported worker group" + ) def supports(self, profile: ExecutionProfile, worker_count: int) -> bool: """Return whether profile and worker count are declared as executable.""" @@ -718,6 +1301,11 @@ def to_dict(self) -> dict[str, Any]: "max_retries": self.policy.max_retries, "exactness": self.policy.exactness.value, } + elif isinstance(self.policy, TorchPolicy): + policy = { + "kind": "torch", + **self.policy.to_dict(), + } else: policy = { "kind": "framework_native", @@ -858,6 +1446,174 @@ def from_dict(cls, value: Mapping[str, Any]) -> DistributionSpec: ), exactness=DistributedExactness(policy_value["exactness"]), ) + elif kind == "torch": + execution_plan = _mapping( + policy_value["execution_plan"], "Torch execution_plan" + ) + stage_values = _sequence( + execution_plan["stages"], "Torch execution stages" + ) + stages = tuple( + TorchStageSpec( + stage_id=_string( + _mapping(item, "Torch stage")["stage_id"], + "Torch stage_id", + ), + worker_loop_ref=_string( + _mapping(item, "Torch stage")["worker_loop_ref"], + "Torch worker_loop_ref", + ), + input_roles=tuple( + _string(role, "Torch input role") + for role in _sequence( + _mapping(item, "Torch stage")["input_roles"], + "Torch input_roles", + ) + ), + depends_on=tuple( + _string(dep, "Torch dependency") + for dep in _sequence( + _mapping(item, "Torch stage").get("depends_on", ()), + "Torch depends_on", + ) + ), + checkpoint_from_stage=_mapping(item, "Torch stage").get( + "checkpoint_from_stage" + ), + metric_mapping=_mapping( + _mapping(item, "Torch stage").get("metric_mapping", {}), + "Torch metric_mapping", + ), + checkpoint_required=_boolean( + _mapping(item, "Torch stage").get( + "checkpoint_required", True + ), + "Torch checkpoint_required", + ), + ) + for item in stage_values + ) + plan_value: TorchExecutionPlan + if execution_plan.get("kind") == "single": + plan_value = SingleStageTorchPlan( + api_version=_positive_integer( + execution_plan["api_version"], + "Torch execution plan api_version", + ), + stage=stages[0], + ) + else: + plan_value = ComponentStageTorchPlan( + api_version=_positive_integer( + execution_plan["api_version"], + "Torch execution plan api_version", + ), + stages=stages, + final_stage_id=_string( + execution_plan["final_stage_id"], + "Torch final_stage_id", + ), + ) + routes = tuple( + TorchDatasetRoute( + role=_string( + _mapping(item, "Torch route")["role"], "Torch role" + ), + mode=_string( + _mapping(item, "Torch route")["mode"], "Torch mode" + ), + required=_boolean( + _mapping(item, "Torch route").get("required", True), + "Torch route required", + ), + min_total_rows_if_present=_non_negative_integer( + _mapping(item, "Torch route").get( + "min_total_rows_if_present", 1 + ), + "Torch minimum total rows", + ), + min_rows_per_worker=_non_negative_integer( + _mapping(item, "Torch route").get("min_rows_per_worker", 1), + "Torch minimum rows per worker", + ), + empty_rank_policy=_string( + _mapping(item, "Torch route").get( + "empty_rank_policy", "reject" + ), + "Torch empty rank policy", + ), + max_rows=_mapping(item, "Torch route").get("max_rows"), + max_bytes_per_worker=_mapping(item, "Torch route").get( + "max_bytes_per_worker" + ), + ) + for item in _sequence( + policy_value["dataset_routing"], "Torch dataset_routing" + ) + ) + policy = TorchPolicy( + torch_runtime_api_version=_positive_integer( + policy_value["torch_runtime_api_version"], + "Torch Runtime API version", + ), + loop_owner=_string(policy_value["loop_owner"], "Torch loop_owner"), + parallelism_id=_string( + policy_value["parallelism_id"], "Torch parallelism_id" + ), + dataset_routing=routes, + execution_plan=plan_value, + state_layout=_string( + policy_value["state_layout"], "Torch state_layout" + ), + metric_reducers={ + _string(name, "Torch metric name"): MetricReduction(value) + for name, value in _mapping( + policy_value["metric_reducers"], "Torch metric_reducers" + ).items() + }, + backend=_string( + policy_value.get("backend", "auto"), "Torch backend" + ), + checkpoint_owner_rank=_non_negative_integer( + policy_value.get("checkpoint_owner_rank", 0), + "Torch checkpoint_owner_rank", + ), + resume_supported=_boolean( + policy_value.get("resume_supported", True), + "Torch resume_supported", + ), + same_world_size_resume=policy_value.get("same_world_size_resume"), + rank_seeded=_boolean( + policy_value.get("rank_seeded", True), "Torch rank_seeded" + ), + checkpoint_adapter_ref=policy_value.get("checkpoint_adapter_ref"), + evidence_adapter_ref=policy_value.get("evidence_adapter_ref"), + global_loss_reducer_ref=policy_value.get("global_loss_reducer_ref"), + global_loss_reducer_api_version=policy_value.get( + "global_loss_reducer_api_version" + ), + global_loss_reducer_code_digest=policy_value.get( + "global_loss_reducer_code_digest" + ), + composite_loss_schema_id=policy_value.get( + "composite_loss_schema_id" + ), + capabilities=tuple( + _string(capability, "Torch capability") + for capability in _sequence( + policy_value.get("capabilities", ()), "Torch capabilities" + ) + ), + max_replicated_bytes_per_worker=( + _positive_integer( + policy_value["max_replicated_bytes_per_worker"], + "Torch max_replicated_bytes_per_worker", + ) + if policy_value.get("max_replicated_bytes_per_worker") + is not None + else None + ), + ) elif kind == "framework_native": policy = FrameworkNativePolicy( framework=_string(policy_value["framework"], "framework"), @@ -953,6 +1709,7 @@ def digest(self) -> str: __all__ = [ + "ComponentStageTorchPlan", "CollectivePolicy", "DistributedExactness", "DistributionSpec", @@ -968,6 +1725,11 @@ def digest(self) -> str: "ResultPolicy", "StateCoordination", "StateField", + "SingleStageTorchPlan", + "TorchDatasetRoute", + "TorchExecutionPlan", + "TorchPolicy", + "TorchStageSpec", "WorkerRange", "WorkerResources", ] diff --git a/src/tributo/algorithms/api/execution.py b/src/tributo/algorithms/api/execution.py index c2a1ee1..d5dca01 100644 --- a/src/tributo/algorithms/api/execution.py +++ b/src/tributo/algorithms/api/execution.py @@ -8,6 +8,7 @@ import re from collections.abc import Mapping from dataclasses import dataclass, field +from types import MappingProxyType from typing import Any, cast from tributo._common.immutable import FrozenDict @@ -23,6 +24,10 @@ FORMAL_DISTRIBUTED_STRATEGY_CONTRACTS, AlgorithmRequest, ) +from tributo.algorithms.api.torch_runtime import ( + TorchRecoveryEnvelope, + TorchStageRunIdentity, +) from tributo.util.annotations import PublicAPI _DIGEST = re.compile(r"^[0-9a-f]{64}$") @@ -80,6 +85,7 @@ class ExecutionRequest: worker_count: int resources_per_worker: WorkerResources | None = None resume_from: str | None = None + torch_recovery: TorchRecoveryEnvelope | None = None def __post_init__(self) -> None: if not isinstance(self.algorithm_request, AlgorithmRequest): @@ -107,6 +113,16 @@ def __post_init__(self) -> None: ) if self.resume_from is not None: _non_empty(self.resume_from, "resume_from") + if self.torch_recovery is not None and not isinstance( + self.torch_recovery, TorchRecoveryEnvelope + ): + raise AlgorithmConfigurationError( + "torch_recovery must be a TorchRecoveryEnvelope" + ) + if self.torch_recovery is not None and self.resume_from is not None: + raise AlgorithmConfigurationError( + "resume_from and torch_recovery are mutually exclusive" + ) @PublicAPI(stability="alpha") @@ -352,6 +368,475 @@ def from_dict(cls, value: Mapping[str, object]) -> StateCoordinationEvidence: ) from exc +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchRoleExecutionEvidence: + """Per-role coverage evidence for a Torch Stage.""" + + role: str + mode: str + required: bool + present: bool + empty_rank_policy: str + expected_rows: int | None + observed_rows: int + rows_per_rank: tuple[int, ...] + replicated_bytes_per_worker: int | None = None + binding_digest: str | None = None + + def __post_init__(self) -> None: + _non_empty(self.role, "Torch role") + if self.mode not in {"split_exact", "replicate", "split_framework"}: + raise AlgorithmConfigurationError("invalid Torch role evidence mode") + if not isinstance(self.required, bool) or not isinstance(self.present, bool): + raise AlgorithmConfigurationError( + "Torch role evidence flags must be boolean" + ) + if self.empty_rank_policy not in {"reject", "zero_contribution"}: + raise AlgorithmConfigurationError("invalid Torch role empty-rank policy") + for name, value in ( + ("expected_rows", self.expected_rows), + ("observed_rows", self.observed_rows), + ("replicated_bytes_per_worker", self.replicated_bytes_per_worker), + ): + if value is not None and ( + not isinstance(value, int) or isinstance(value, bool) or value < 0 + ): + raise AlgorithmConfigurationError( + f"Torch role evidence {name} must be non-negative" + ) + rows = tuple(self.rows_per_rank) + if any( + not isinstance(value, int) or isinstance(value, bool) or value < 0 + for value in rows + ): + raise AlgorithmConfigurationError("Torch rows_per_rank is malformed") + if self.present: + if self.mode == "replicate": + if not rows or len(set(rows)) != 1 or rows[0] != self.observed_rows: + raise AlgorithmConfigurationError( + "Torch replicated role evidence is not identical" + ) + elif sum(rows) != self.observed_rows: + raise AlgorithmConfigurationError("Torch role evidence rows do not sum") + if not self.present and (self.observed_rows != 0 or any(rows)): + raise AlgorithmConfigurationError( + "absent Torch role must have zero evidence" + ) + if ( + self.binding_digest is not None + and _DIGEST.fullmatch(self.binding_digest) is None + ): + raise AlgorithmConfigurationError("Torch role binding_digest is invalid") + object.__setattr__(self, "rows_per_rank", rows) + + def to_dict(self) -> dict[str, Any]: + return { + "role": self.role, + "mode": self.mode, + "required": self.required, + "present": self.present, + "empty_rank_policy": self.empty_rank_policy, + "expected_rows": self.expected_rows, + "observed_rows": self.observed_rows, + "rows_per_rank": list(self.rows_per_rank), + "replicated_bytes_per_worker": self.replicated_bytes_per_worker, + "binding_digest": self.binding_digest, + } + + @classmethod + def from_dict(cls, value: Mapping[str, object]) -> "TorchRoleExecutionEvidence": + try: + return cls( + role=cast(str, value["role"]), + mode=cast(str, value["mode"]), + required=cast(bool, value["required"]), + present=cast(bool, value["present"]), + empty_rank_policy=cast(str, value["empty_rank_policy"]), + expected_rows=cast(int | None, value.get("expected_rows")), + observed_rows=cast(int, value["observed_rows"]), + rows_per_rank=tuple( + cast(tuple[int, ...], value.get("rows_per_rank", ())) + ), + replicated_bytes_per_worker=cast( + int | None, value.get("replicated_bytes_per_worker") + ), + binding_digest=cast(str | None, value.get("binding_digest")), + ) + except KeyError as exc: + raise AlgorithmConfigurationError( + "Torch role evidence is missing fields" + ) from exc + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class ReplicatedTorchStateEvidence: + """Evidence that all workers expose one synchronized model.""" + + model_digests_by_rank: Mapping[int, str] + global_model_digest: str + + def __post_init__(self) -> None: + records = dict(self.model_digests_by_rank) + if not records or any( + not isinstance(rank, int) or rank < 0 for rank in records + ): + raise AlgorithmConfigurationError( + "replicated Torch rank evidence is required" + ) + if any(_DIGEST.fullmatch(value) is None for value in records.values()): + raise AlgorithmConfigurationError( + "replicated Torch model digest is invalid" + ) + if _DIGEST.fullmatch(self.global_model_digest or "") is None: + raise AlgorithmConfigurationError( + "replicated Torch global digest is invalid" + ) + if set(records.values()) != {self.global_model_digest}: + raise AlgorithmConfigurationError( + "replicated Torch ranks are not synchronized" + ) + object.__setattr__( + self, + "model_digests_by_rank", + MappingProxyType(dict(sorted(records.items()))), + ) + + def to_dict(self) -> dict[str, Any]: + return { + "model_digests_by_rank": { + str(k): v for k, v in self.model_digests_by_rank.items() + }, + "global_model_digest": self.global_model_digest, + } + + @classmethod + def from_dict(cls, value: Mapping[str, object]) -> "ReplicatedTorchStateEvidence": + raw = _mapping(value.get("model_digests_by_rank", {}), "model_digests_by_rank") + return cls( + model_digests_by_rank={ + int(rank): cast(str, digest) for rank, digest in raw.items() + }, + global_model_digest=cast(str, value["global_model_digest"]), + ) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class ComponentStageEvidence: + """Evidence for one Core-owned component Stage.""" + + stage_id: str + workers: tuple[WorkerExecutionEvidence, ...] + roles: tuple[TorchRoleExecutionEvidence, ...] + state_digest: str + checkpoint_descriptor_digest: str | None = None + + def __post_init__(self) -> None: + _non_empty(self.stage_id, "Torch stage evidence stage_id") + workers = tuple(self.workers) + if not workers or any( + not isinstance(item, WorkerExecutionEvidence) for item in workers + ): + raise AlgorithmConfigurationError("Torch Stage evidence requires Workers") + roles = tuple(self.roles) + if any(not isinstance(item, TorchRoleExecutionEvidence) for item in roles): + raise AlgorithmConfigurationError( + "Torch Stage evidence roles are malformed" + ) + for name, value in (("state_digest", self.state_digest),): + if _DIGEST.fullmatch(value or "") is None: + raise AlgorithmConfigurationError( + f"Torch Stage evidence {name} is invalid" + ) + if ( + self.checkpoint_descriptor_digest is not None + and _DIGEST.fullmatch(self.checkpoint_descriptor_digest) is None + ): + raise AlgorithmConfigurationError( + "Torch Stage evidence checkpoint_descriptor_digest is invalid" + ) + object.__setattr__(self, "workers", workers) + object.__setattr__(self, "roles", roles) + + def to_dict(self) -> dict[str, Any]: + return { + "stage_id": self.stage_id, + "workers": [item.to_dict() for item in self.workers], + "roles": [item.to_dict() for item in self.roles], + "state_digest": self.state_digest, + "checkpoint_descriptor_digest": self.checkpoint_descriptor_digest, + } + + @classmethod + def from_dict(cls, value: Mapping[str, object]) -> "ComponentStageEvidence": + workers_value = value.get("workers", ()) + roles_value = value.get("roles", ()) + if not isinstance(workers_value, (list, tuple)) or not isinstance( + roles_value, (list, tuple) + ): + raise AlgorithmConfigurationError( + "Torch Stage evidence worker/role fields are invalid" + ) + return cls( + stage_id=cast(str, value["stage_id"]), + workers=tuple( + WorkerExecutionEvidence.from_dict(item) for item in workers_value + ), + roles=tuple( + TorchRoleExecutionEvidence.from_dict(item) for item in roles_value + ), + state_digest=cast(str, value["state_digest"]), + checkpoint_descriptor_digest=cast( + str | None, value.get("checkpoint_descriptor_digest") + ), + ) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchExecutionEvidence: + """Layout-aware execution evidence attached only to Torch receipts.""" + + identity: TorchStageRunIdentity + run_config_name: str + policy_digest: str + parallelism_id: str + state_layout: str + workers: tuple[WorkerExecutionEvidence, ...] + roles: tuple[TorchRoleExecutionEvidence, ...] + replicated_state: ReplicatedTorchStateEvidence | None = None + stages: tuple[ComponentStageEvidence, ...] = () + composition_digest: str | None = None + final_stage_id: str | None = None + reducer_id: str | None = None + reducer_api_version: int | None = None + reducer_schema_id: str | None = None + reducer_code_digest: str | None = None + torch_runtime_api_version: int = 1 + reducer_branch: str | None = None + reducer_evidence: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.identity, TorchStageRunIdentity): + raise AlgorithmConfigurationError("Torch execution identity is required") + if not isinstance(self.run_config_name, str) or not self.run_config_name: + raise AlgorithmConfigurationError("Torch run_config_name is required") + if self.run_config_name != self.identity.run_config_name: + raise AlgorithmConfigurationError("Torch RunConfig name drifted") + if ( + self.torch_runtime_api_version != 1 + or self.identity.torch_runtime_api_version != 1 + ): + raise AlgorithmConfigurationError( + "Torch runtime API version must be exactly 1" + ) + if _DIGEST.fullmatch(self.policy_digest or "") is None: + raise AlgorithmConfigurationError("Torch policy digest is invalid") + if self.identity.policy_digest != self.policy_digest: + raise AlgorithmConfigurationError("Torch execution policy digest drifted") + if self.identity.execution_plan_digest == "": + raise AlgorithmConfigurationError("Torch execution plan digest is required") + if not isinstance(self.parallelism_id, str) or not self.parallelism_id: + raise AlgorithmConfigurationError("Torch parallelism_id is required") + if self.state_layout not in {"replicated", "component", "sharded"}: + raise AlgorithmConfigurationError("Torch state layout is invalid") + workers = tuple(self.workers) + roles = tuple(self.roles) + if not workers or any( + not isinstance(item, WorkerExecutionEvidence) for item in workers + ): + raise AlgorithmConfigurationError( + "Torch execution evidence requires workers" + ) + if any(not isinstance(item, TorchRoleExecutionEvidence) for item in roles): + raise AlgorithmConfigurationError( + "Torch execution role evidence is malformed" + ) + if self.state_layout == "replicated": + if self.replicated_state is None or self.stages: + raise AlgorithmConfigurationError( + "replicated Torch evidence has invalid state payload" + ) + elif self.state_layout == "component": + if not self.stages or self.replicated_state is not None: + raise AlgorithmConfigurationError( + "component Torch evidence has invalid stage payload" + ) + if _DIGEST.fullmatch(self.composition_digest or "") is None: + raise AlgorithmConfigurationError( + "component Torch composition digest is required" + ) + if len({stage.stage_id for stage in self.stages}) != len(self.stages): + raise AlgorithmConfigurationError( + "component Torch Stage IDs must be unique" + ) + expected_composition = hashlib.sha256( + json.dumps( + [stage.to_dict() for stage in self.stages], + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + ).hexdigest() + if self.composition_digest != expected_composition: + raise AlgorithmConfigurationError( + "component Torch composition digest drifted" + ) + if self.final_stage_id not in {stage.stage_id for stage in self.stages}: + raise AlgorithmConfigurationError( + "component Torch final_stage_id is invalid" + ) + else: + if not self.stages and self.replicated_state is None: + raise AlgorithmConfigurationError( + "sharded Torch evidence requires state payload" + ) + if any(item.required and not item.present for item in roles): + raise AlgorithmConfigurationError("required Torch role evidence is absent") + for stage in self.stages: + stage_workers = tuple(stage.workers) + ranks = tuple(sorted(item.rank for item in stage_workers)) + if ranks != tuple(range(len(stage_workers))) or any( + item.world_size != len(stage_workers) for item in stage_workers + ): + raise AlgorithmConfigurationError( + "Torch Stage worker evidence is incomplete" + ) + reducer_fields = ( + self.reducer_id, + self.reducer_api_version, + self.reducer_schema_id, + self.reducer_code_digest, + ) + if any(value is not None for value in reducer_fields) and not all( + value is not None for value in reducer_fields + ): + raise AlgorithmConfigurationError("Torch reducer evidence is incomplete") + if self.reducer_api_version is not None and self.reducer_api_version != 1: + raise AlgorithmConfigurationError( + "Torch reducer API version must be exactly 1" + ) + if self.reducer_branch is not None and ( + not isinstance(self.reducer_branch, str) or not self.reducer_branch + ): + raise AlgorithmConfigurationError("Torch reducer branch must be non-empty") + evidence = dict(self.reducer_evidence) + if any(not isinstance(name, str) or not name for name in evidence): + raise AlgorithmConfigurationError( + "Torch reducer evidence names are invalid" + ) + if any( + not isinstance(value, (str, int, float, bool, type(None))) + for value in evidence.values() + ): + raise AlgorithmConfigurationError( + "Torch reducer evidence must be JSON scalar" + ) + if any( + isinstance(value, float) and not math.isfinite(value) + for value in evidence.values() + ): + raise AlgorithmConfigurationError("Torch reducer evidence must be finite") + object.__setattr__(self, "reducer_evidence", MappingProxyType(evidence)) + for name, value in (("reducer_code_digest", self.reducer_code_digest),): + if value is not None and _DIGEST.fullmatch(value) is None: + raise AlgorithmConfigurationError(f"Torch {name} is invalid") + object.__setattr__(self, "workers", workers) + object.__setattr__(self, "roles", roles) + + def to_dict(self) -> dict[str, Any]: + return { + "identity": self.identity.to_dict(), + "run_config_name": self.run_config_name, + "policy_digest": self.policy_digest, + "parallelism_id": self.parallelism_id, + "state_layout": self.state_layout, + "workers": [item.to_dict() for item in self.workers], + "roles": [item.to_dict() for item in self.roles], + "replicated_state": self.replicated_state.to_dict() + if self.replicated_state + else None, + "stages": [item.to_dict() for item in self.stages], + "composition_digest": self.composition_digest, + "final_stage_id": self.final_stage_id, + "reducer_id": self.reducer_id, + "reducer_api_version": self.reducer_api_version, + "reducer_schema_id": self.reducer_schema_id, + "reducer_code_digest": self.reducer_code_digest, + "torch_runtime_api_version": self.torch_runtime_api_version, + "reducer_branch": self.reducer_branch, + "reducer_evidence": dict(self.reducer_evidence), + } + + @classmethod + def from_dict(cls, value: Mapping[str, object]) -> "TorchExecutionEvidence": + identity = TorchStageRunIdentity.from_dict( + _mapping(value.get("identity", {}), "Torch evidence identity") + ) + workers_value = value.get("workers", ()) + roles_value = value.get("roles", ()) + if not isinstance(workers_value, (list, tuple)) or not isinstance( + roles_value, (list, tuple) + ): + raise AlgorithmConfigurationError( + "Torch execution evidence worker/role fields are invalid" + ) + workers = tuple( + WorkerExecutionEvidence.from_dict(item) + for item in workers_value + if isinstance(item, Mapping) + ) + if len(workers) != len(workers_value): + raise AlgorithmConfigurationError( + "Torch worker evidence entries are invalid" + ) + roles = tuple( + TorchRoleExecutionEvidence.from_dict(item) + for item in roles_value + if isinstance(item, Mapping) + ) + if len(roles) != len(roles_value): + raise AlgorithmConfigurationError("Torch role evidence entries are invalid") + stages_value = value.get("stages", ()) + if not isinstance(stages_value, (list, tuple)): + raise AlgorithmConfigurationError( + "Torch Stage evidence entries are invalid" + ) + replicated_value = value.get("replicated_state") + replicated = ( + ReplicatedTorchStateEvidence.from_dict(replicated_value) + if isinstance(replicated_value, Mapping) + else None + ) + return cls( + identity=identity, + run_config_name=cast(str, value["run_config_name"]), + policy_digest=cast(str, value["policy_digest"]), + parallelism_id=cast(str, value["parallelism_id"]), + state_layout=cast(str, value["state_layout"]), + workers=workers, + roles=roles, + replicated_state=replicated, + stages=tuple( + ComponentStageEvidence.from_dict(item) for item in stages_value + ), + composition_digest=cast(str | None, value.get("composition_digest")), + final_stage_id=cast(str | None, value.get("final_stage_id")), + reducer_id=cast(str | None, value.get("reducer_id")), + reducer_api_version=cast(int | None, value.get("reducer_api_version")), + reducer_schema_id=cast(str | None, value.get("reducer_schema_id")), + reducer_code_digest=cast(str | None, value.get("reducer_code_digest")), + torch_runtime_api_version=cast( + int, value.get("torch_runtime_api_version", 1) + ), + reducer_branch=cast(str | None, value.get("reducer_branch")), + reducer_evidence=cast( + Mapping[str, object], value.get("reducer_evidence", {}) + ), + ) + + @PublicAPI(stability="alpha") @dataclass(frozen=True) class ExecutionReceipt: @@ -376,6 +861,7 @@ class ExecutionReceipt: runtime_owned: bool = False resource_preflight: str = "validated" api_version: int = 1 + torch_evidence: TorchExecutionEvidence | None = None def __post_init__(self) -> None: _non_empty(self.run_id, "run_id") @@ -396,6 +882,15 @@ def __post_init__(self) -> None: object.__setattr__(self, "profile", profile) object.__setattr__(self, "strategy", strategy) object.__setattr__(self, "result_policy", result_policy) + if strategy is DistributionStrategy.RAY_TRAIN_TORCH: + if not isinstance(self.torch_evidence, TorchExecutionEvidence): + raise AlgorithmConfigurationError( + "RAY_TRAIN_TORCH receipts require TorchExecutionEvidence" + ) + elif self.torch_evidence is not None: + raise AlgorithmConfigurationError( + "non-Torch receipts must not contain TorchExecutionEvidence" + ) if ( not isinstance(self.api_version, int) or isinstance(self.api_version, bool) @@ -723,6 +1218,9 @@ def to_dict(self) -> dict[str, Any]: "cluster_resources": dict(sorted(self.cluster_resources.items())), "runtime_owned": self.runtime_owned, "resource_preflight": self.resource_preflight, + "torch_evidence": self.torch_evidence.to_dict() + if self.torch_evidence is not None + else None, "distributed": self.distributed, "cross_node": self.cross_node, "cluster_distributed": self.cluster_distributed, @@ -731,8 +1229,12 @@ def to_dict(self) -> dict[str, Any]: __all__ = [ + "ComponentStageEvidence", "ExecutionReceipt", "ExecutionRequest", + "ReplicatedTorchStateEvidence", "StateCoordinationEvidence", + "TorchExecutionEvidence", + "TorchRoleExecutionEvidence", "WorkerExecutionEvidence", ] diff --git a/src/tributo/algorithms/api/models.py b/src/tributo/algorithms/api/models.py index 76f6044..36f32c7 100644 --- a/src/tributo/algorithms/api/models.py +++ b/src/tributo/algorithms/api/models.py @@ -146,7 +146,7 @@ class ExecutionMode(str, Enum): JOBLIB_ESTIMATOR = "joblib_estimator" PARALLEL_ENSEMBLE = "parallel_ensemble" ITERATIVE_OPTIMIZATION = "iterative_optimization" - TRAINING_RECIPE_V2 = "training_recipe_v2" + RAY_TRAIN_TORCH = "ray_train_torch" @PublicAPI(stability="alpha") @@ -162,7 +162,7 @@ class RuntimeTopology(str, Enum): RAY_JOBLIB_ESTIMATOR = "ray_joblib_estimator" RAY_PARALLEL_ENSEMBLE = "ray_parallel_ensemble" RAY_ITERATIVE_OPTIMIZATION = "ray_iterative_optimization" - RAY_TRAIN_RECIPE_V2 = "ray_train_recipe_v2" + RAY_TRAIN_TORCH = "ray_train_torch" @dataclass(frozen=True) @@ -250,12 +250,12 @@ class FormalDistributedStrategyContract: ), ) ), - DistributionStrategy.RAY_TRAIN_RECIPE_V2: FormalDistributedStrategyContract( - execution_mode=ExecutionMode.TRAINING_RECIPE_V2, - runtime_id="tributo.ray_train_recipe_v2", - topology=RuntimeTopology.RAY_TRAIN_RECIPE_V2, - input_distribution=InputDistribution.SHARDED, - state_coordination=StateCoordination.ALL_REDUCE, + DistributionStrategy.RAY_TRAIN_TORCH: FormalDistributedStrategyContract( + execution_mode=ExecutionMode.RAY_TRAIN_TORCH, + runtime_id="tributo.ray_train_torch", + topology=RuntimeTopology.RAY_TRAIN_TORCH, + input_distribution=InputDistribution.ROLE_ROUTED, + state_coordination=StateCoordination.TORCH_MANAGED, worker_input_adapter_ref=( "tributo.integrations.algorithm_inputs.ingestion:" "prepare_ray_train_input" @@ -609,7 +609,7 @@ def __post_init__(self) -> None: ExecutionMode.JOBLIB_ESTIMATOR, ExecutionMode.PARALLEL_ENSEMBLE, ExecutionMode.ITERATIVE_OPTIMIZATION, - ExecutionMode.TRAINING_RECIPE_V2, + ExecutionMode.RAY_TRAIN_TORCH, } if formal_mode and ( self.runtime_id is None or self.worker_input_adapter_ref is None @@ -689,6 +689,7 @@ class RuntimeBinding: strategy: DistributionStrategy | None = None distribution_digest: str | None = None resume_from: str | None = None + torch_recovery: Mapping[str, Any] | None = None memory_bytes: int | None = None def __post_init__(self) -> None: @@ -800,7 +801,7 @@ def __post_init__(self) -> None: RuntimeTopology.RAY_JOBLIB_ESTIMATOR, RuntimeTopology.RAY_PARALLEL_ENSEMBLE, RuntimeTopology.RAY_ITERATIVE_OPTIMIZATION, - RuntimeTopology.RAY_TRAIN_RECIPE_V2, + RuntimeTopology.RAY_TRAIN_TORCH, } if topology in formal_topologies and ( self.framework_parallelism != 1 or self.result_reducer_ref is not None @@ -849,6 +850,24 @@ def __post_init__(self) -> None: raise AlgorithmConfigurationError( "legacy RuntimeBinding must not carry formal resume state" ) + if self.torch_recovery is not None: + if self.strategy is not DistributionStrategy.RAY_TRAIN_TORCH: + raise AlgorithmConfigurationError( + "torch_recovery requires the Ray Train Torch runtime" + ) + from tributo.algorithms.api.torch_runtime import TorchRecoveryEnvelope + + try: + recovery = TorchRecoveryEnvelope.from_dict(self.torch_recovery) + except (TypeError, ValueError) as exc: + raise AlgorithmConfigurationError( + "runtime torch_recovery is malformed" + ) from exc + object.__setattr__(self, "torch_recovery", deep_freeze(recovery.to_dict())) + if self.resume_from is not None and self.torch_recovery is not None: + raise AlgorithmConfigurationError( + "runtime resume_from and torch_recovery are mutually exclusive" + ) @PublicAPI(stability="alpha") @@ -1706,6 +1725,9 @@ def to_dict(self, *, include_plan_id: bool = True) -> dict[str, Any]: ), "distribution_digest": self.runtime.distribution_digest, "resume_from": self.runtime.resume_from, + "torch_recovery": deep_thaw(self.runtime.torch_recovery) + if self.runtime.torch_recovery is not None + else None, } if self.runtime.memory_bytes is not None: runtime_payload["memory_bytes"] = self.runtime.memory_bytes diff --git a/src/tributo/algorithms/api/support.py b/src/tributo/algorithms/api/support.py index 9497863..614cef7 100644 --- a/src/tributo/algorithms/api/support.py +++ b/src/tributo/algorithms/api/support.py @@ -17,6 +17,7 @@ from tributo.algorithms.api.distribution import ( DistributionStrategy, ExecutionProfile, + TorchPolicy, ) from tributo.algorithms.api.errors import AlgorithmConfigurationError from tributo.util.annotations import PublicAPI @@ -118,6 +119,8 @@ class AlgorithmSupportEvidence: expires_at: datetime | None = None revoked_at: datetime | None = None revocation_reason: str | None = None + torch_runtime_api_version: int | None = None + torch_policy_digest: str | None = None def __post_init__(self) -> None: for name in ( @@ -167,6 +170,23 @@ def __post_init__(self) -> None: _digest(self.wheel_sha256, "wheel_sha256") for contract_digest in contract_digests: _digest(contract_digest, "contract_digest") + if self.torch_runtime_api_version is not None: + if ( + not isinstance(self.torch_runtime_api_version, int) + or isinstance(self.torch_runtime_api_version, bool) + or self.torch_runtime_api_version != 1 + ): + raise AlgorithmConfigurationError( + "torch_runtime_api_version must be exactly 1 when provided" + ) + if self.torch_policy_digest is not None: + _digest(self.torch_policy_digest, "torch_policy_digest") + if (self.torch_runtime_api_version is None) != ( + self.torch_policy_digest is None + ): + raise AlgorithmConfigurationError( + "Torch support fields must be supplied together" + ) for name in ("issued_at", "expires_at", "revoked_at"): timestamp = getattr(self, name) if timestamp is not None and ( @@ -207,6 +227,8 @@ def evidence_id(self) -> str: "issued_at": self.issued_at.isoformat(), "gate": self.gate, "result_reference": self.result_reference, + "torch_runtime_api_version": self.torch_runtime_api_version, + "torch_policy_digest": self.torch_policy_digest, } return hashlib.sha256( json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() @@ -230,6 +252,7 @@ def matches( """Match every immutable package, contract, profile, and semantic key.""" distribution = descriptor.registration.distribution_spec environment = descriptor.registration.environment + policy = distribution.policy if distribution is not None else None ray_requirements = [ requirement for requirement in environment.dependencies @@ -260,6 +283,19 @@ def matches( in SpecifierSet(descriptor.tributo_version_spec) and Version(self.python_version) in SpecifierSet(environment.python) and ray_matches + and ( + ( + not isinstance(policy, TorchPolicy) + and self.torch_runtime_api_version is None + and self.torch_policy_digest is None + ) + or ( + isinstance(policy, TorchPolicy) + and self.torch_runtime_api_version + == policy.torch_runtime_api_version + and self.torch_policy_digest == policy.digest + ) + ) ) diff --git a/src/tributo/algorithms/api/torch_runtime.py b/src/tributo/algorithms/api/torch_runtime.py new file mode 100644 index 0000000..09b2a5f --- /dev/null +++ b/src/tributo/algorithms/api/torch_runtime.py @@ -0,0 +1,2246 @@ +"""Public, framework-neutral contracts for the Ray Train Torch runtime. + +The module deliberately avoids importing Torch, Ray, or any algorithm package at +import time. Runtime implementations live behind lazy integration modules and +consume these immutable values through the public API. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import re +import threading +from collections.abc import Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Generator, Protocol, cast, runtime_checkable +from urllib.parse import urlsplit + +from tributo._common.immutable import deep_freeze +from tributo.algorithms.api.errors import ( + AlgorithmConfigurationError, + AlgorithmExecutionError, +) +from tributo.util.annotations import PublicAPI + +if TYPE_CHECKING: + from tributo.algorithms.spi.execution import RuntimeExecutionEnvelope + from tributo.algorithms.spi.torch import TorchStageContext + + +def _canonical_json(value: Mapping[str, Any]) -> str: + try: + return json.dumps( + value, + allow_nan=False, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ) + except (TypeError, ValueError) as exc: + raise AlgorithmConfigurationError( + "Torch runtime metadata must be canonical JSON" + ) from exc + + +def _digest(value: Mapping[str, Any]) -> str: + return hashlib.sha256(_canonical_json(value).encode("utf-8")).hexdigest() + + +def _finite_number(value: object, field_name: str) -> float: + if ( + not isinstance(value, (int, float)) + or isinstance(value, bool) + or not math.isfinite(float(value)) + ): + raise AlgorithmConfigurationError(f"{field_name} must be finite") + return float(value) + + +def _validate_scalar_numerator(value: object, field_name: str) -> None: + """Validate a differentiable zero-dimensional tensor without importing Torch.""" + ndim = getattr(value, "ndim", None) + if not isinstance(ndim, int) or isinstance(ndim, bool) or ndim != 0: + raise AlgorithmConfigurationError( + f"{field_name} must be a differentiable zero-dimensional Tensor" + ) + detach = getattr(value, "detach", None) + item = getattr(value, "item", None) + if not callable(detach) or not callable(item): + raise AlgorithmConfigurationError( + f"{field_name} must provide callable detach() and item()" + ) + try: + detached = detach() + detached_item = getattr(detached, "item", None) + if not callable(detached_item) or not math.isfinite(float(detached_item())): + raise AlgorithmConfigurationError(f"{field_name} must be finite") + except (TypeError, ValueError) as exc: + raise AlgorithmConfigurationError(f"{field_name} must be finite") from exc + + +def _digest_value(value: object, field_name: str) -> str: + if not isinstance(value, str) or len(value) != 64: + raise AlgorithmConfigurationError(f"{field_name} must be a SHA-256 digest") + try: + int(value, 16) + except ValueError as exc: + raise AlgorithmConfigurationError( + f"{field_name} must be a lower-case hexadecimal digest" + ) from exc + if value != value.lower(): + raise AlgorithmConfigurationError(f"{field_name} must be lower-case") + return value + + +def _validate_bounded_evidence( + value: object, + *, + path: str = "evidence", + depth: int = 0, +) -> None: + if depth > 4: + raise AlgorithmConfigurationError("Torch evidence nesting is too deep") + if isinstance(value, Mapping): + if len(value) > 64: + raise AlgorithmConfigurationError("Torch evidence has too many fields") + for key, nested in value.items(): + if ( + not isinstance(key, str) + or not key + or len(key) > 128 + or re.fullmatch(r"[a-z][a-z0-9_.-]*", key) is None + or key.casefold() + in { + "path", + "uri", + "locator", + "credential", + "credentials", + "secret", + "secrets", + } + or key.casefold().endswith(("_path", "_uri", "_locator")) + ): + raise AlgorithmConfigurationError( + f"Torch evidence key {path}.{key!r} is invalid" + ) + _validate_bounded_evidence(nested, path=f"{path}.{key}", depth=depth + 1) + return + if isinstance(value, (list, tuple)): + if len(value) > 64: + raise AlgorithmConfigurationError( + f"Torch evidence list {path} is too large" + ) + for index, nested in enumerate(value): + _validate_bounded_evidence(nested, path=f"{path}[{index}]", depth=depth + 1) + return + if isinstance(value, str): + if ( + len(value) > 1024 + or value.startswith(("/", "~/")) + or re.match(r"^[A-Za-z][A-Za-z0-9+.-]*://", value) + ): + raise AlgorithmConfigurationError(f"Torch evidence value {path} is unsafe") + return + if value is None or isinstance(value, bool) or isinstance(value, int): + return + if isinstance(value, float) and math.isfinite(value): + return + raise AlgorithmConfigurationError(f"Torch evidence value {path} is not portable") + + +_RUN_TOKEN = re.compile(r"^[a-f0-9]{8,128}$") + + +def _run_component(value: str, field_name: str) -> str: + """Return a canonical path-safe run component.""" + if not isinstance(value, str) or not value or _RUN_TOKEN.fullmatch(value) is None: + raise AlgorithmConfigurationError( + f"{field_name} must be canonical lower-case UUID/hex" + ) + return value + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchStageRunIdentity: + """Identity shared by one logical Torch Stage Run and its checkpoints.""" + + run_id: str + invocation_id: str + stage_id: str + torch_runtime_api_version: int + algorithm: str + implementation_id: str + implementation_code_digest: str + policy_digest: str + execution_plan_digest: str + plan_digest: str | None = None + + def __post_init__(self) -> None: + for name in ( + "run_id", + "invocation_id", + "stage_id", + "algorithm", + "implementation_id", + ): + value = getattr(self, name) + if not isinstance(value, str) or not value: + raise AlgorithmConfigurationError(f"{name} must be non-empty") + if self.torch_runtime_api_version != 1: + raise AlgorithmConfigurationError( + "torch_runtime_api_version must be exactly 1" + ) + _digest_value(self.implementation_code_digest, "implementation_code_digest") + _digest_value(self.policy_digest, "policy_digest") + _digest_value(self.execution_plan_digest, "execution_plan_digest") + if self.plan_digest is not None: + _digest_value(self.plan_digest, "plan_digest") + + def to_dict(self) -> dict[str, Any]: + return { + "run_id": self.run_id, + "invocation_id": self.invocation_id, + "stage_id": self.stage_id, + "torch_runtime_api_version": self.torch_runtime_api_version, + "algorithm": self.algorithm, + "implementation_id": self.implementation_id, + "implementation_code_digest": self.implementation_code_digest, + "policy_digest": self.policy_digest, + "execution_plan_digest": self.execution_plan_digest, + "plan_digest": self.plan_digest, + } + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "TorchStageRunIdentity": + try: + return cls(**dict(value)) + except (TypeError, KeyError) as exc: + raise AlgorithmConfigurationError( + "invalid Torch Stage Run identity" + ) from exc + + @property + def identity_digest(self) -> str: + return _digest(self.to_dict()) + + @property + def run_config_name(self) -> str: + return torch_run_config_name(self) + + +@PublicAPI(stability="alpha") +def torch_run_config_name(identity: TorchStageRunIdentity) -> str: + """Build the Core-owned deterministic Ray ``RunConfig.name``.""" + if not isinstance(identity, TorchStageRunIdentity): + raise AlgorithmConfigurationError("Torch run identity is required") + run_id = _run_component(identity.run_id, "run_id") + invocation_id = _run_component(identity.invocation_id, "invocation_id") + stage_digest = hashlib.sha256(identity.stage_id.encode("utf-8")).hexdigest()[:16] + name = ( + f"tributo-torch-v1-{run_id}-{invocation_id}-" + f"{stage_digest}-{identity.identity_digest}" + ) + if len(name) > 255 or any( + char not in "abcdefghijklmnopqrstuvwxyz0123456789-_" for char in name + ): + raise AlgorithmConfigurationError("Torch RunConfig.name is not filesystem-safe") + return name + + +@PublicAPI(stability="alpha") +def claim_torch_run_directory( + storage_path: str | os.PathLike[str], + identity: TorchStageRunIdentity, + *, + status: str = "running", +) -> Path: + """Atomically claim a Stage directory and its identity manifest. + + Existing directories are reusable only for the same identity and a + retryable status. Nothing is overwritten or silently renamed. + """ + if status not in {"created", "running", "failed_retryable"}: + raise AlgorithmConfigurationError("Torch run directory status is not retryable") + root = Path(storage_path) + raw_storage = str(storage_path) + if "://" in raw_storage and not raw_storage.startswith("file://"): + # Use the filesystem abstraction only for the identity manifest. The + # actual Ray checkpoint payload remains owned by Ray's storage backend. + try: + import pyarrow.fs as pafs + + filesystem, prefix = pafs.FileSystem.from_uri(raw_storage) + run_path = f"{prefix.rstrip('/')}/{torch_run_config_name(identity)}" + manifest_path = f"{run_path}/torch_identity_manifest.json" + filesystem.create_dir(run_path, recursive=True) + info = filesystem.get_file_info(manifest_path) + payload = { + "schema_version": 1, + "snapshot_schema": 1, + "identity": identity.to_dict(), + "status": status, + "stale": False, + } + if info.type is pafs.FileType.File: + existing = json.loads( + filesystem.open_input_file(manifest_path).read().decode("utf-8") + ) + if ( + not isinstance(existing, Mapping) + or existing.get("identity") != identity.to_dict() + or existing.get("snapshot_schema", 1) != 1 + or existing.get("stale", False) + or existing.get("status") + not in {"created", "running", "failed_retryable"} + ): + raise AlgorithmExecutionError( + "Torch remote run directory identity collision or stale snapshot" + ) + elif info.type is pafs.FileType.NotFound: + with filesystem.open_output_stream(manifest_path) as stream: + stream.write((_canonical_json(payload) + "\n").encode("utf-8")) + else: + raise AlgorithmExecutionError( + "Torch remote run identity manifest has an invalid type" + ) + return Path(f"{raw_storage.rstrip('/')}/{torch_run_config_name(identity)}") + except AlgorithmExecutionError: + raise + except (ImportError, OSError, TypeError, ValueError) as exc: + raise AlgorithmExecutionError( + "failed to claim remote Torch run identity manifest" + ) from exc + if not root.is_absolute(): + raise AlgorithmConfigurationError("Torch storage_path must be absolute") + directory = root / torch_run_config_name(identity) + manifest = directory / "torch_identity_manifest.json" + try: + directory.mkdir(parents=True, exist_ok=False) + except FileExistsError: + if not manifest.is_file(): + raise AlgorithmExecutionError( + "Torch run directory has no identity manifest" + ) from None + try: + existing = json.loads(manifest.read_text(encoding="utf-8")) + except (OSError, TypeError, ValueError) as exc: + raise AlgorithmExecutionError( + "Torch run identity manifest is damaged" + ) from exc + if ( + not isinstance(existing, Mapping) + or existing.get("identity") != identity.to_dict() + ): + raise AlgorithmExecutionError( + "Torch run directory identity collision" + ) from None + if existing.get("snapshot_schema", 1) != 1 or existing.get("stale", False): + raise AlgorithmExecutionError( + "Torch run directory contains a stale snapshot" + ) from None + if existing.get("status") not in {"created", "running", "failed_retryable"}: + raise AlgorithmExecutionError( + "Torch run directory is terminal or stale" + ) from None + return directory + payload = { + "schema_version": 1, + "snapshot_schema": 1, + "identity": identity.to_dict(), + "status": status, + "stale": False, + } + try: + fd = os.open(manifest, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + with os.fdopen(fd, "w", encoding="utf-8") as stream: + stream.write(_canonical_json(payload)) + stream.write("\n") + except OSError as exc: + raise AlgorithmExecutionError( + "failed to atomically write Torch identity manifest" + ) from exc + return directory + + +@PublicAPI(stability="alpha") +def validate_torch_retry_identity( + descriptor: "TorchCheckpointDescriptor", + identity: TorchStageRunIdentity, + *, + world_size: int, +) -> None: + """Reject stale Ray retry snapshots before model state is deserialized.""" + if not isinstance(descriptor, TorchCheckpointDescriptor): + raise AlgorithmExecutionError("retry checkpoint has no Torch descriptor") + if descriptor.identity.to_dict() != identity.to_dict(): + raise AlgorithmExecutionError( + "retry checkpoint identity does not match current Stage" + ) + if descriptor.world_size != world_size: + raise AlgorithmExecutionError( + "retry checkpoint world size does not match current Stage" + ) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchLossContribution: + """One ordinary differentiable loss numerator and its denominator.""" + + numerator: object + normalizer: float + + def __post_init__(self) -> None: + _validate_scalar_numerator(self.numerator, "loss numerator") + normalizer = _finite_number(self.normalizer, "loss normalizer") + if normalizer < 0: + raise AlgorithmConfigurationError("loss normalizer must be non-negative") + object.__setattr__(self, "normalizer", normalizer) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchCompositeLossContribution: + """Named differentiable components reduced by an algorithm-owned reducer.""" + + schema_id: str + differentiable_components: Mapping[str, object] + normalizer_components: Mapping[str, float] + evidence: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.schema_id, str) or not self.schema_id: + raise AlgorithmConfigurationError("composite loss schema_id is required") + components: dict[str, object] = dict(self.differentiable_components) + normalizers: dict[str, float] = dict(self.normalizer_components) + if not components: + raise AlgorithmConfigurationError( + "composite loss requires differentiable components" + ) + if not normalizers: + raise AlgorithmConfigurationError( + "composite loss requires normalizer components" + ) + for name, value in normalizers.items(): + if not isinstance(name, str) or not name: + raise AlgorithmConfigurationError( + "composite loss normalizers must be named and non-negative" + ) + normalized_value = _finite_number(value, f"normalizer[{name}]") + if normalized_value < 0: + raise AlgorithmConfigurationError( + "composite loss normalizers must be named and non-negative" + ) + normalizers[name] = normalized_value + for component_name, component_value in components.items(): + if ( + not isinstance(component_name, str) + or not component_name + or component_value is None + ): + raise AlgorithmConfigurationError( + "differentiable component names and values are required" + ) + _validate_scalar_numerator(component_value, f"component[{component_name}]") + object.__setattr__( + self, "differentiable_components", MappingProxyType(components) + ) + object.__setattr__(self, "normalizer_components", MappingProxyType(normalizers)) + _validate_bounded_evidence(self.evidence) + object.__setattr__(self, "evidence", deep_freeze(self.evidence)) + + +TorchStepLoss = TorchLossContribution | TorchCompositeLossContribution + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchMetricContribution: + """A metric numerator/denominator pair reduced by Core.""" + + numerator: float + normalizer: float + + def __post_init__(self) -> None: + numerator = _finite_number(self.numerator, "metric numerator") + normalizer = _finite_number(self.normalizer, "metric normalizer") + if normalizer < 0: + raise AlgorithmConfigurationError("metric normalizer must be non-negative") + if normalizer == 0 and numerator != 0: + raise AlgorithmConfigurationError( + "zero metric normalizer requires a zero numerator" + ) + object.__setattr__(self, "numerator", numerator) + object.__setattr__(self, "normalizer", normalizer) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchAccumulationWindow: + """State describing one optimizer accumulation window.""" + + index: int + expected_micro_batches: int + observed_micro_batches: int = 0 + normalizer_total: float = 0.0 + + def __post_init__(self) -> None: + if ( + not isinstance(self.index, int) + or isinstance(self.index, bool) + or not isinstance(self.expected_micro_batches, int) + or isinstance(self.expected_micro_batches, bool) + or not isinstance(self.observed_micro_batches, int) + or isinstance(self.observed_micro_batches, bool) + or self.index < 0 + or self.expected_micro_batches < 1 + ): + raise AlgorithmConfigurationError("invalid accumulation window") + if not 0 <= self.observed_micro_batches <= self.expected_micro_batches: + raise AlgorithmConfigurationError("invalid observed micro-batch count") + total = _finite_number(self.normalizer_total, "accumulation normalizer total") + if total < 0: + raise AlgorithmConfigurationError( + "accumulation normalizer total must be non-negative" + ) + object.__setattr__(self, "normalizer_total", total) + + def add(self, normalizer: float) -> "TorchAccumulationWindow": + """Return the next immutable window state after one micro-batch.""" + value = _finite_number(normalizer, "loss normalizer") + if value < 0: + raise AlgorithmConfigurationError("loss normalizer must be non-negative") + if self.observed_micro_batches >= self.expected_micro_batches: + raise AlgorithmExecutionError("accumulation window is already complete") + return TorchAccumulationWindow( + index=self.index, + expected_micro_batches=self.expected_micro_batches, + observed_micro_batches=self.observed_micro_batches + 1, + normalizer_total=self.normalizer_total + value, + ) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchBackwardContext: + """Public callbacks supplied by the Core Torch loop to the backward helper.""" + + world_size: int + backward: Any + reduce_normalizer: Any + finalize_window: Any + reduce_window_normalizer: Any | None = None + compose_composite: Any | None = None + + def __post_init__(self) -> None: + if self.world_size < 1 or not callable(self.backward): + raise AlgorithmConfigurationError("invalid Torch backward context") + if not callable(self.reduce_normalizer) or not callable(self.finalize_window): + raise AlgorithmConfigurationError( + "Torch backward callbacks must be callable" + ) + if self.reduce_window_normalizer is not None and not callable( + self.reduce_window_normalizer + ): + raise AlgorithmConfigurationError( + "Torch window normalizer callback must be callable" + ) + if self.compose_composite is not None and not callable(self.compose_composite): + raise AlgorithmConfigurationError( + "Torch composite backward callback must be callable" + ) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchBackwardResult: + """Result returned after one loss contribution is submitted.""" + + local_normalizer: float + global_normalizer: float + window_complete: bool + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchMetricPolicy: + """Explicit metric reducer names keyed by metric identity.""" + + reducers: Mapping[str, str] + + def __post_init__(self) -> None: + normalized = dict(self.reducers) + if any(not isinstance(k, str) or not k for k in normalized): + raise AlgorithmConfigurationError("metric reducer names are required") + object.__setattr__(self, "reducers", MappingProxyType(normalized)) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchMetricReductionContext: + """Core callbacks used by the metric reduction helper.""" + + reduce: Any + + def __post_init__(self) -> None: + if not callable(self.reduce): + raise AlgorithmConfigurationError("metric reduction callback is required") + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchMetricReductionResult: + """Reduced metric values and their evidence.""" + + values: Mapping[str, float] + evidence: Mapping[str, object] = field(default_factory=dict) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchGlobalLossContext: + """Context passed to an algorithm-owned global loss reducer.""" + + world_size: int + policy_digest: str + execution_plan_digest: str + config: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + if ( + not isinstance(self.world_size, int) + or isinstance(self.world_size, bool) + or self.world_size < 1 + ): + raise AlgorithmConfigurationError( + "Torch global loss world_size must be positive" + ) + _digest_value(self.policy_digest, "policy_digest") + _digest_value(self.execution_plan_digest, "execution_plan_digest") + object.__setattr__(self, "config", deep_freeze(self.config)) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchCompositeGlobalState: + """Names-preserving detached component state produced by Core AllReduce.""" + + components: Mapping[str, float] + normalizers: Mapping[str, float] + + def __post_init__(self) -> None: + components = dict(self.components) + normalizers = dict(self.normalizers) + if not components or not normalizers: + raise AlgorithmConfigurationError( + "global component and normalizer state are required" + ) + for name, value in components.items(): + if not isinstance(name, str) or not name: + raise AlgorithmConfigurationError("global component names are required") + components[name] = _finite_number(value, f"global component[{name}]") + for name, value in normalizers.items(): + if not isinstance(name, str) or not name: + raise AlgorithmConfigurationError( + "global normalizer names are required" + ) + normalizers[name] = _finite_number(value, f"global normalizer[{name}]") + if value < 0: + raise AlgorithmConfigurationError( + f"global normalizer[{name}] must be non-negative" + ) + object.__setattr__(self, "components", MappingProxyType(components)) + object.__setattr__(self, "normalizers", MappingProxyType(normalizers)) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchGlobalLossReduction: + """Deterministic reducer output shared by all ranks.""" + + status: str + coefficients: Mapping[str, float] = field(default_factory=dict) + branch: str | None = None + evidence: Mapping[str, object] = field(default_factory=dict) + failure_code: str | None = None + metrics: Mapping[str, TorchMetricContribution] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.status not in {"accepted", "rejected"}: + raise AlgorithmConfigurationError("loss reduction status is invalid") + for name in ("branch", "failure_code"): + value = getattr(self, name) + if value is not None and ( + not isinstance(value, str) + or not value + or len(value) > 128 + or re.fullmatch(r"[a-z][a-z0-9_.-]*", value) is None + ): + raise AlgorithmConfigurationError( + f"loss reduction {name} must be a bounded namespaced value" + ) + coefficients = dict(self.coefficients) + evidence = dict(self.evidence) + metrics = dict(self.metrics) + if self.status == "accepted": + if not coefficients or self.failure_code is not None: + raise AlgorithmConfigurationError( + "accepted loss reduction requires coefficients" + ) + for name, value in coefficients.items(): + coefficients[name] = _finite_number(value, f"coefficient[{name}]") + elif coefficients or not self.failure_code: + raise AlgorithmConfigurationError( + "rejected loss reduction requires a failure code and no coefficients" + ) + if any( + not isinstance(value, TorchMetricContribution) for value in metrics.values() + ): + raise AlgorithmConfigurationError("loss reduction metrics are invalid") + if len(evidence) > 64 or any( + not isinstance(name, str) + or not name + or len(name) > 128 + or re.fullmatch(r"[a-z][a-z0-9_.-]*", name) is None + or not isinstance(value, (str, int, float, bool, type(None))) + or (isinstance(value, str) and len(value) > 1024) + or (isinstance(value, float) and not math.isfinite(value)) + for name, value in evidence.items() + ): + raise AlgorithmConfigurationError( + "loss reduction evidence must be bounded JSON scalars" + ) + object.__setattr__(self, "coefficients", MappingProxyType(coefficients)) + object.__setattr__(self, "evidence", MappingProxyType(evidence)) + object.__setattr__(self, "metrics", MappingProxyType(metrics)) + + def to_dict(self) -> dict[str, Any]: + return { + "status": self.status, + "coefficients": dict(self.coefficients), + "branch": self.branch, + "evidence": dict(self.evidence), + "failure_code": self.failure_code, + "metrics": { + name: { + "numerator": value.numerator, + "normalizer": value.normalizer, + } + for name, value in self.metrics.items() + }, + } + + +@runtime_checkable +@PublicAPI(stability="alpha") +class TorchGlobalLossReducer(Protocol): + """Algorithm-owned deterministic global loss reducer.""" + + api_version: int + reducer_id: str + component_schema_id: str + code_digest: str + + def reduce( + self, + config: Mapping[str, object], + global_state: TorchCompositeGlobalState, + context: TorchGlobalLossContext, + ) -> TorchGlobalLossReduction: ... + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchPreflightTokenData: + """Immutable identity data produced by Torch preflight.""" + + run_id: str + invocation_id: str + algorithm: str + implementation_ref: str + implementation_code_digest: str + policy_digest: str + execution_plan_digest: str + runtime_id: str + reducer_identity: str | None = None + plan_digest: str | None = None + + def __post_init__(self) -> None: + for name in ( + "run_id", + "invocation_id", + "algorithm", + "implementation_ref", + "runtime_id", + ): + if not isinstance(getattr(self, name), str) or not getattr(self, name): + raise AlgorithmConfigurationError(f"{name} must be non-empty") + _digest_value(self.implementation_code_digest, "implementation_code_digest") + _digest_value(self.policy_digest, "policy_digest") + _digest_value(self.execution_plan_digest, "execution_plan_digest") + if self.reducer_identity is not None and ( + not isinstance(self.reducer_identity, str) + or ":" not in self.reducer_identity + ): + raise AlgorithmConfigurationError( + "reducer_identity must be qualified when provided" + ) + if self.plan_digest is not None: + _digest_value(self.plan_digest, "plan_digest") + + def to_dict(self) -> dict[str, Any]: + return { + "run_id": self.run_id, + "invocation_id": self.invocation_id, + "algorithm": self.algorithm, + "implementation_ref": self.implementation_ref, + "implementation_code_digest": self.implementation_code_digest, + "policy_digest": self.policy_digest, + "execution_plan_digest": self.execution_plan_digest, + "runtime_id": self.runtime_id, + "reducer_identity": self.reducer_identity, + "plan_digest": self.plan_digest, + } + + +@PublicAPI(stability="alpha") +@runtime_checkable +class TorchCheckpointPayloadDraft(Protocol): + """Optional Core payload draft hook used by checkpoint report tests.""" + + checkpoint_dir: str | os.PathLike[str] + + def report( + self, + *, + metrics: Mapping[str, object], + stage_context: object, + completed_step: int, + ) -> None: ... + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchCheckpointRef: + """Driver-owned reference to an actual Ray/framework checkpoint.""" + + checkpoint: object + descriptor_digest: str | None = None + source_stage_id: str | None = None + descriptor: "TorchCheckpointDescriptor | None" = None + + def __post_init__(self) -> None: + if self.checkpoint is None: + raise AlgorithmConfigurationError("Torch checkpoint reference is required") + if self.descriptor_digest is not None: + _digest_value(self.descriptor_digest, "descriptor_digest") + if self.descriptor is not None: + if not isinstance(self.descriptor, TorchCheckpointDescriptor): + raise AlgorithmConfigurationError( + "Torch checkpoint descriptor is invalid" + ) + if ( + self.descriptor_digest is not None + and self.descriptor.digest != self.descriptor_digest + ): + raise AlgorithmConfigurationError( + "Torch checkpoint descriptor digest mismatch" + ) + + def __getstate__(self) -> object: + raise TypeError("TorchCheckpointRef is driver-local and cannot be serialized") + + def __copy__(self) -> object: + raise TypeError("TorchCheckpointRef cannot be copied") + + def __deepcopy__(self, memo: dict[int, object]) -> object: + del memo + raise TypeError("TorchCheckpointRef cannot be copied") + + def close(self) -> None: + """Release a driver-side checkpoint handle when the backend supports it.""" + closer = getattr(self.checkpoint, "close", None) + if callable(closer): + closer() + + def __enter__(self) -> "TorchCheckpointRef": + return self + + def __exit__(self, exc_type: object, exc: object, traceback: object) -> None: + del exc_type, exc, traceback + self.close() + + +@PublicAPI(stability="alpha") +class TorchPreflightLease: + """Invocation-local one-shot ownership token for a preflight result.""" + + __slots__ = ("_data", "_state", "_lock") + + def __init__(self, data: TorchPreflightTokenData) -> None: + if not isinstance(data, TorchPreflightTokenData): + raise AlgorithmConfigurationError("preflight lease requires TokenData") + self._data = data + self._state = "fresh" + self._lock = threading.Lock() + + @property + def state(self) -> str: + with self._lock: + return self._state + + def _matches( + self, *, run_id: str, invocation_id: str, plan_digest: str, runtime_id: str + ) -> bool: + return ( + self._data.run_id == run_id + and self._data.invocation_id == invocation_id + and ( + self._data.plan_digest is None + and self._data.execution_plan_digest == plan_digest + or self._data.plan_digest is not None + and self._data.plan_digest == plan_digest + ) + and self._data.runtime_id == runtime_id + ) + + def claim( + self, + *, + run_id: str, + invocation_id: str, + plan_digest: str, + runtime_id: str, + ) -> None: + with self._lock: + if self._state != "fresh": + raise AlgorithmExecutionError("Torch preflight lease is not fresh") + if not self._matches( + run_id=run_id, + invocation_id=invocation_id, + plan_digest=plan_digest, + runtime_id=runtime_id, + ): + raise AlgorithmExecutionError("Torch preflight lease identity mismatch") + self._state = "claimed" + + def consume( + self, + *, + run_id: str, + invocation_id: str, + plan_digest: str, + runtime_id: str, + ) -> TorchPreflightTokenData: + with self._lock: + if self._state != "claimed": + raise AlgorithmExecutionError("Torch preflight lease is not claimed") + if not self._matches( + run_id=run_id, + invocation_id=invocation_id, + plan_digest=plan_digest, + runtime_id=runtime_id, + ): + raise AlgorithmExecutionError("Torch preflight token identity mismatch") + self._state = "consumed" + return self._data + + def close(self) -> None: + with self._lock: + if self._state in {"fresh", "claimed"}: + self._state = "closed" + + @property + def data(self) -> TorchPreflightTokenData: + with self._lock: + if self._state not in {"fresh", "claimed"}: + raise AlgorithmExecutionError("Torch preflight lease is closed") + return self._data + + def __copy__(self) -> object: + raise TypeError("TorchPreflightLease cannot be copied") + + def __deepcopy__(self, memo: dict[int, object]) -> object: + del memo + raise TypeError("TorchPreflightLease cannot be copied") + + def __getstate__(self) -> object: + raise TypeError("TorchPreflightLease cannot be serialized") + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchWorkerControlEnvelope: + """Credential-free serialized control for initial Stage Checkpoint input.""" + + schema_version: int + run_id: str + invocation_id: str + source_stage_id: str | None + target_stage_id: str + purpose: str + checkpoint_locator: "TorchCheckpointLocator" + checkpoint_descriptor_digest: str + policy_digest: str + execution_plan_digest: str + + def __post_init__(self) -> None: + if self.schema_version != 1: + raise AlgorithmConfigurationError( + "unsupported Torch worker control version" + ) + if self.purpose not in {"stage_dependency", "cross_run_initial_recovery"}: + raise AlgorithmConfigurationError("invalid Torch worker control purpose") + if self.purpose == "stage_dependency" and ( + not isinstance(self.source_stage_id, str) or not self.source_stage_id + ): + raise AlgorithmConfigurationError( + "stage_dependency control requires a source_stage_id" + ) + for name in ("run_id", "invocation_id", "target_stage_id"): + if not isinstance(getattr(self, name), str) or not getattr(self, name): + raise AlgorithmConfigurationError(f"{name} must be non-empty") + if self.source_stage_id == self.target_stage_id: + raise AlgorithmConfigurationError( + "Torch control source and target Stage must differ" + ) + if not isinstance(self.checkpoint_locator, TorchCheckpointLocator): + raise AlgorithmConfigurationError("Torch checkpoint locator is invalid") + _digest_value(self.checkpoint_descriptor_digest, "checkpoint_descriptor_digest") + if ( + self.checkpoint_locator.descriptor_digest + != self.checkpoint_descriptor_digest + ): + raise AlgorithmConfigurationError( + "Torch control locator and descriptor digests do not match" + ) + _digest_value(self.policy_digest, "policy_digest") + _digest_value(self.execution_plan_digest, "execution_plan_digest") + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "run_id": self.run_id, + "invocation_id": self.invocation_id, + "source_stage_id": self.source_stage_id, + "target_stage_id": self.target_stage_id, + "purpose": self.purpose, + "checkpoint_locator": self.checkpoint_locator.to_dict(), + "checkpoint_descriptor_digest": self.checkpoint_descriptor_digest, + "policy_digest": self.policy_digest, + "execution_plan_digest": self.execution_plan_digest, + } + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "TorchWorkerControlEnvelope": + raw_locator = value.get("checkpoint_locator") + if not isinstance(raw_locator, Mapping): + raise AlgorithmConfigurationError( + "Torch worker control locator must be a typed locator" + ) + locator = TorchCheckpointLocator.from_dict(raw_locator) + try: + return cls( + schema_version=value["schema_version"], + run_id=value["run_id"], + invocation_id=value["invocation_id"], + source_stage_id=value.get("source_stage_id"), + target_stage_id=value["target_stage_id"], + purpose=value["purpose"], + checkpoint_locator=locator, + checkpoint_descriptor_digest=value["checkpoint_descriptor_digest"], + policy_digest=value["policy_digest"], + execution_plan_digest=value["execution_plan_digest"], + ) + except KeyError as exc: + raise AlgorithmConfigurationError( + "Torch worker control is missing a field" + ) from exc + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchRuntimeExecutionEnvelope: + """Driver-local Torch envelope carrying a claimed preflight lease.""" + + base: "RuntimeExecutionEnvelope" + preflight_lease: TorchPreflightLease + + def __post_init__(self) -> None: + if not hasattr(self.base, "plan") or not hasattr(self.base, "run_id"): + raise AlgorithmConfigurationError( + "Torch envelope requires a RuntimeExecutionEnvelope" + ) + if not isinstance(self.preflight_lease, TorchPreflightLease): + raise AlgorithmConfigurationError( + "Torch envelope requires a preflight lease" + ) + + def __getstate__(self) -> object: + raise TypeError("Torch runtime envelope cannot be serialized") + + def __copy__(self) -> object: + raise TypeError("Torch runtime envelope cannot be copied") + + def __deepcopy__(self, memo: dict[int, object]) -> object: + del memo + raise TypeError("Torch runtime envelope cannot be copied") + + +def _loss_numerator_and_normalizer(loss: TorchStepLoss) -> tuple[object, float]: + if isinstance(loss, TorchLossContribution): + return loss.numerator, loss.normalizer + if isinstance(loss, TorchCompositeLossContribution): + raise AlgorithmConfigurationError( + "composite loss requires invoke_torch_global_loss_reducer" + ) + raise AlgorithmConfigurationError("unsupported Torch loss contribution") + + +@PublicAPI(stability="alpha") +def apply_torch_loss_backward( + loss: TorchStepLoss, + window: TorchAccumulationWindow, + context: TorchBackwardContext, +) -> TorchBackwardResult: + """Submit one ordinary loss contribution to the Core accumulation helper.""" + if isinstance(loss, TorchCompositeLossContribution): + if context.compose_composite is None: + raise AlgorithmConfigurationError( + "composite loss requires a Core composite backward callback" + ) + numerator = context.compose_composite(loss) + if not isinstance(getattr(numerator, "ndim", None), int): + raise AlgorithmExecutionError( + "composite backward callback returned invalid numerator" + ) + normalizer = sum(loss.normalizer_components.values()) + else: + numerator, normalizer = _loss_numerator_and_normalizer(loss) + context.backward(numerator) + next_window = window.add(normalizer) + complete = next_window.observed_micro_batches >= next_window.expected_micro_batches + if isinstance(loss, TorchCompositeLossContribution): + if complete: + # The reducer coefficients already encode the global objective; + # applying the ordinary numerator/normalizer scale would normalize + # it a second time. ``compose_composite`` must include the + # Core-required world-size factor before returning its scalar. + context.finalize_window(1.0) + return TorchBackwardResult( + normalizer, + next_window.normalizer_total, + complete, + ) + if not complete: + return TorchBackwardResult(normalizer, next_window.normalizer_total, False) + reduce = context.reduce_window_normalizer or context.reduce_normalizer + global_normalizer = _finite_number( + reduce(next_window.normalizer_total), "global loss normalizer" + ) + if global_normalizer <= 0: + raise AlgorithmExecutionError("global loss normalizer must be positive") + context.finalize_window(context.world_size / global_normalizer) + return TorchBackwardResult(normalizer, global_normalizer, True) + + +@PublicAPI(stability="alpha") +def reduce_torch_metrics( + contributions: Mapping[str, TorchMetricContribution], + policy: TorchMetricPolicy, + context: TorchMetricReductionContext, +) -> TorchMetricReductionResult: + """Reduce explicit metric numerator/normalizer contributions.""" + if set(contributions) != set(policy.reducers): + raise AlgorithmConfigurationError( + "metric contribution names do not match policy" + ) + values: dict[str, float] = {} + for name, contribution in contributions.items(): + if not isinstance(contribution, TorchMetricContribution): + raise AlgorithmConfigurationError(f"metric {name!r} is not typed") + result = context.reduce(name, contribution, policy.reducers[name]) + values[name] = _finite_number(result, f"metric[{name}]") + return TorchMetricReductionResult(values) + + +@PublicAPI(stability="alpha") +def invoke_torch_global_loss_reducer( + contribution: TorchCompositeLossContribution, + global_state: TorchCompositeGlobalState, + reducer: TorchGlobalLossReducer, + context: TorchGlobalLossContext, +) -> TorchGlobalLossReduction: + """Invoke an algorithm-owned reducer without granting collective access.""" + if not isinstance(contribution, TorchCompositeLossContribution): + raise AlgorithmConfigurationError("composite loss contribution is required") + if contribution.schema_id != reducer.component_schema_id: + raise AlgorithmConfigurationError( + "composite loss schema does not match reducer" + ) + if getattr(reducer, "api_version", None) != 1: + raise AlgorithmConfigurationError( + "Torch global loss reducer API version must be 1" + ) + if ( + not isinstance(getattr(reducer, "reducer_id", None), str) + or not reducer.reducer_id + ): + raise AlgorithmConfigurationError( + "Torch global loss reducer identity is required" + ) + _digest_value(getattr(reducer, "code_digest", None), "reducer code digest") + result = reducer.reduce(context.config, global_state, context) + if not isinstance(result, TorchGlobalLossReduction): + raise AlgorithmExecutionError("global loss reducer returned an invalid result") + if result.status == "accepted" and set(result.coefficients) != set( + contribution.differentiable_components + ): + raise AlgorithmExecutionError( + "Torch global loss reducer coefficients do not match components" + ) + return result + + +@PublicAPI(stability="alpha") +def report_torch_checkpoint( + metrics: Mapping[str, object], + payload_draft: TorchCheckpointPayloadDraft, + stage_context: TorchStageContext, + completed_step: int, +) -> None: + """Report a Core-validated Torch checkpoint through Ray Train.""" + if completed_step < 0: + raise AlgorithmConfigurationError("completed_step must be non-negative") + if not isinstance(metrics, Mapping): + raise AlgorithmConfigurationError("Torch checkpoint metrics must be a mapping") + + def validate_metric_metadata(value: object, path: str = "metrics") -> None: + if isinstance(value, Mapping): + for raw_key, nested in value.items(): + if not isinstance(raw_key, str): + raise AlgorithmConfigurationError( + "Torch checkpoint metric keys must be strings" + ) + key = raw_key.casefold() + if key in { + "path", + "uri", + "locator", + "checkpoint", + "credential", + "credentials", + "secret", + "secrets", + "password", + "token", + } or key.endswith(("_path", "_uri", "_locator")): + raise AlgorithmConfigurationError( + f"Torch checkpoint metrics contain a Core-owned field: {path}.{raw_key}" + ) + validate_metric_metadata(nested, f"{path}.{raw_key}") + elif isinstance(value, (list, tuple)): + for index, nested in enumerate(value): + validate_metric_metadata(nested, f"{path}[{index}]") + elif isinstance(value, str): + if value.startswith(("/", "~/")) or re.match( + r"^[A-Za-z][A-Za-z0-9+.-]*://", value + ): + raise AlgorithmConfigurationError( + f"Torch checkpoint metrics contain a path or URI value: {path}" + ) + + validate_metric_metadata(metrics) + reporter = getattr(payload_draft, "report", None) + checkpoint_dir = getattr(payload_draft, "checkpoint_dir", None) + if not callable(reporter) or checkpoint_dir is None: + raise AlgorithmExecutionError( + "Torch checkpoints require a Core payload draft with checkpoint_dir and report()" + ) + root = Path(checkpoint_dir) + if root.is_symlink() or not root.is_dir(): + raise AlgorithmExecutionError("Torch checkpoint payload directory is missing") + root = root.resolve() + runtime = getattr(stage_context, "runtime", None) + identity = getattr(runtime, "run_identity", None) + if identity is None: + raise AlgorithmExecutionError( + "Torch checkpoint Stage context has no Run identity" + ) + if runtime is None: + raise AlgorithmExecutionError("Torch checkpoint Stage context has no runtime") + binding_digest = getattr(runtime, "input_binding_digest", None) + if not isinstance(binding_digest, str) or len(binding_digest) != 64: + raise AlgorithmExecutionError( + "Torch checkpoint Stage context has no complete input binding digest" + ) + descriptor_path = root / "torch_checkpoint_descriptor.json" + if descriptor_path.is_symlink(): + raise AlgorithmExecutionError( + "Torch checkpoint descriptor must not be a symlink" + ) + if descriptor_path.is_file(): + try: + previous = TorchCheckpointDescriptor.from_dict( + json.loads(descriptor_path.read_text(encoding="utf-8")) + ) + except (OSError, TypeError, ValueError) as exc: + raise AlgorithmExecutionError( + "existing Torch checkpoint descriptor is malformed" + ) from exc + if previous.identity != identity: + raise AlgorithmExecutionError( + "Torch checkpoint identity changed within a Stage" + ) + if completed_step <= previous.completed_step: + raise AlgorithmExecutionError( + "Torch checkpoint completed_step must increase monotonically" + ) + evidence_path = root / "torch_execution_evidence.json" + if evidence_path.is_symlink(): + raise AlgorithmExecutionError( + "Torch checkpoint execution evidence must not be a symlink" + ) + evidence_payload = { + name: metrics[name] + for name in ( + "execution_workers", + "model_state_digest", + "reducer_id", + "reducer_api_version", + "reducer_schema_id", + "reducer_code_digest", + "reducer_branch", + "reducer_evidence", + ) + if name in metrics + } + if evidence_payload: + temporary_evidence = root / ".torch_execution_evidence.tmp" + if temporary_evidence.exists() or temporary_evidence.is_symlink(): + raise AlgorithmExecutionError( + "Torch checkpoint evidence temporary path already exists" + ) + try: + fd = os.open( + temporary_evidence, + os.O_WRONLY | os.O_CREAT | os.O_EXCL, + 0o600, + ) + with os.fdopen(fd, "w", encoding="utf-8") as stream: + stream.write(_canonical_json(evidence_payload)) + stream.write("\n") + os.replace(temporary_evidence, evidence_path) + except OSError as exc: + try: + temporary_evidence.unlink(missing_ok=True) + except OSError: + pass + raise AlgorithmExecutionError( + "failed to atomically write Torch checkpoint evidence" + ) from exc + files = _scan_checkpoint_files(root) + if not files: + raise AlgorithmExecutionError("Torch checkpoint payload is empty") + descriptor = TorchCheckpointDescriptor( + schema_version=1, + identity=identity, + run_config_name=identity.run_config_name, + state_layout=getattr(runtime, "state_layout", "replicated"), + world_size=runtime.world_size, + completed_step=completed_step, + policy_digest=runtime.policy_digest, + execution_plan_digest=runtime.execution_plan_digest, + input_binding_digest=binding_digest, + implementation_code_digest=identity.implementation_code_digest, + payload_files=dict(files), + adapter_identity=getattr(runtime, "adapter_identity", None), + resume_supported=getattr(runtime, "resume_supported", True), + same_world_size_resume=getattr(runtime, "same_world_size_resume", True), + ) + temporary_descriptor = root / ".torch_checkpoint_descriptor.tmp" + if temporary_descriptor.exists() or temporary_descriptor.is_symlink(): + raise AlgorithmExecutionError( + "Torch checkpoint descriptor temporary path already exists" + ) + try: + fd = os.open( + temporary_descriptor, + os.O_WRONLY | os.O_CREAT | os.O_EXCL, + 0o600, + ) + with os.fdopen(fd, "w", encoding="utf-8") as stream: + stream.write(_canonical_json(descriptor.to_dict())) + stream.write("\n") + os.replace(temporary_descriptor, descriptor_path) + except OSError as exc: + try: + temporary_descriptor.unlink(missing_ok=True) + except OSError: + pass + raise AlgorithmExecutionError( + "failed to atomically write Torch checkpoint descriptor" + ) from exc + report_metrics = dict(metrics) + report_metrics["checkpoint_descriptor"] = descriptor.to_dict() + reporter( + metrics=report_metrics, + stage_context=stage_context, + completed_step=completed_step, + ) + + +def _scan_checkpoint_files(root: Path) -> dict[str, str]: + resolved_root = root.resolve() + files: dict[str, str] = {} + for path in sorted(root.rglob("*")): + if path.is_symlink() or not path.resolve().is_relative_to(resolved_root): + raise AlgorithmExecutionError("Torch checkpoint payload escapes its root") + if path.name in { + "torch_checkpoint_descriptor.json", + "torch_stage_commit.json", + ".metadata.json", + }: + continue + if path.is_file(): + digest = hashlib.sha256(path.read_bytes()).hexdigest() + files[str(path.relative_to(root))] = digest + return files + + +@PublicAPI(stability="alpha") +def describe_torch_checkpoint( + checkpoint_ref: TorchCheckpointRef, + checkpoint_context: object, +) -> TorchCheckpointDescriptor: + """Read and validate the Core descriptor embedded in a Driver Checkpoint.""" + if not isinstance(checkpoint_ref, TorchCheckpointRef): + raise AlgorithmConfigurationError("Torch checkpoint reference is required") + descriptor = checkpoint_ref.descriptor + with _opened_checkpoint(checkpoint_ref.checkpoint) as root: + descriptor_path = root / "torch_checkpoint_descriptor.json" + if not descriptor_path.is_file() or descriptor_path.is_symlink(): + raise AlgorithmExecutionError("Torch checkpoint has no Core descriptor") + try: + parsed = TorchCheckpointDescriptor.from_dict( + json.loads(descriptor_path.read_text(encoding="utf-8")) + ) + except (OSError, TypeError, ValueError) as exc: + raise AlgorithmExecutionError( + "Torch checkpoint descriptor is malformed" + ) from exc + if descriptor is not None and descriptor.digest != parsed.digest: + raise AlgorithmExecutionError("Torch checkpoint descriptor drifted") + descriptor = parsed + if descriptor.run_config_name != descriptor.identity.run_config_name: + raise AlgorithmExecutionError("Torch checkpoint RunConfig name drifted") + commit_path = root / "torch_stage_commit.json" + if commit_path.exists() or commit_path.is_symlink(): + if commit_path.is_symlink() or not commit_path.is_file(): + raise AlgorithmExecutionError( + "Torch checkpoint commit marker is invalid" + ) + try: + commit = json.loads(commit_path.read_text(encoding="utf-8")) + except (OSError, TypeError, ValueError) as exc: + raise AlgorithmExecutionError( + "Torch checkpoint commit marker is malformed" + ) from exc + if not isinstance(commit, Mapping) or ( + commit.get("identity") != descriptor.identity.to_dict() + or commit.get("descriptor_digest") != descriptor.digest + ): + raise AlgorithmExecutionError( + "Torch checkpoint commit marker does not match descriptor" + ) + files = _scan_checkpoint_files(root) + if dict(descriptor.payload_files) != files: + raise AlgorithmExecutionError( + "Torch checkpoint payload files or digest drifted" + ) + stage = getattr(checkpoint_context, "stage", None) + if ( + stage is not None + and getattr(stage, "stage_id", descriptor.identity.stage_id) + != descriptor.identity.stage_id + ): + raise AlgorithmExecutionError("Torch checkpoint descriptor Stage mismatch") + runtime = getattr(stage, "runtime", None) + expected_identity = getattr(runtime, "run_identity", None) + if expected_identity is not None and descriptor.identity != expected_identity: + raise AlgorithmExecutionError("Torch checkpoint descriptor identity mismatch") + if runtime is not None: + for name in ("policy_digest", "execution_plan_digest"): + expected = getattr(runtime, name, None) + if expected is not None and getattr(descriptor, name) != expected: + raise AlgorithmExecutionError( + f"Torch checkpoint descriptor {name} drifted" + ) + expected_binding = getattr(runtime, "input_binding_digest", None) + if ( + expected_binding is not None + and descriptor.input_binding_digest != expected_binding + ): + raise AlgorithmExecutionError( + "Torch checkpoint descriptor input binding drifted" + ) + expected_adapter = getattr(runtime, "adapter_identity", None) + if descriptor.adapter_identity != expected_adapter: + raise AlgorithmExecutionError( + "Torch checkpoint descriptor Adapter identity drifted" + ) + for name in ("state_layout", "resume_supported", "same_world_size_resume"): + expected = getattr(runtime, name, None) + if getattr(descriptor, name) != expected: + raise AlgorithmExecutionError( + f"Torch checkpoint descriptor {name} drifted" + ) + for name in ("run_id", "invocation_id"): + expected = getattr(checkpoint_context, name, None) + if expected is not None and getattr(descriptor.identity, name) != expected: + raise AlgorithmExecutionError(f"Torch checkpoint descriptor {name} drifted") + return descriptor + + +@contextmanager +def _opened_checkpoint(checkpoint: object) -> Generator[Path, None, None]: + """Open a Ray Checkpoint or local checkpoint directory for validation.""" + if isinstance(checkpoint, (str, Path)): + root = Path(checkpoint) + if not root.is_dir(): + raise AlgorithmExecutionError("Torch checkpoint directory is missing") + yield root + return + as_directory = getattr(checkpoint, "as_directory", None) + if not callable(as_directory): + raise AlgorithmExecutionError("Torch checkpoint cannot be opened") + try: + with as_directory() as directory: + yield Path(directory) + except Exception as exc: + raise AlgorithmExecutionError("Torch checkpoint could not be opened") from exc + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchCheckpointLocator: + """Credential-free persistent location for a validated Torch checkpoint.""" + + uri: str + descriptor_digest: str + schema_version: int = 1 + + def __post_init__(self) -> None: + if self.schema_version != 1: + raise AlgorithmConfigurationError("unsupported Torch locator version") + if not isinstance(self.uri, str) or not self.uri or "\x00" in self.uri: + raise AlgorithmConfigurationError("Torch checkpoint locator URI is invalid") + if self.uri.startswith(("/", "file://")): + raise AlgorithmConfigurationError( + "Torch checkpoint locator must not persist a local path" + ) + try: + parsed = urlsplit(self.uri) + except ValueError as exc: + raise AlgorithmConfigurationError( + "Torch checkpoint locator URI is invalid" + ) from exc + if parsed.username is not None or parsed.password is not None: + raise AlgorithmConfigurationError( + "Torch checkpoint locator must not contain URI userinfo" + ) + if parsed.query or parsed.fragment: + raise AlgorithmConfigurationError( + "Torch checkpoint locator must not contain query or fragment data" + ) + _digest_value(self.descriptor_digest, "descriptor_digest") + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "uri": self.uri, + "descriptor_digest": self.descriptor_digest, + } + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "TorchCheckpointLocator": + try: + return cls( + uri=value["uri"], + descriptor_digest=value["descriptor_digest"], + schema_version=value.get("schema_version", 1), + ) + except KeyError as exc: + raise AlgorithmConfigurationError( + "Torch checkpoint locator is incomplete" + ) from exc + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchRecoveryEnvelope: + """Credential-free cross-Run recovery state for one Torch execution plan.""" + + completed_stage_ids: tuple[str, ...] = () + stage_checkpoints: Mapping[str, TorchCheckpointLocator] = field( + default_factory=dict + ) + active_stage_id: str | None = None + active_checkpoint: TorchCheckpointLocator | None = None + schema_version: int = 1 + + def __post_init__(self) -> None: + if self.schema_version != 1: + raise AlgorithmConfigurationError( + "unsupported Torch recovery envelope version" + ) + completed = tuple(self.completed_stage_ids) + if any(not isinstance(stage, str) or not stage for stage in completed): + raise AlgorithmConfigurationError( + "Torch recovery completed_stage_ids must be non-empty strings" + ) + if len(set(completed)) != len(completed): + raise AlgorithmConfigurationError( + "Torch recovery completed_stage_ids must be unique" + ) + checkpoints: dict[str, TorchCheckpointLocator] = {} + for stage_id, locator in self.stage_checkpoints.items(): + if not isinstance(stage_id, str) or not stage_id: + raise AlgorithmConfigurationError( + "Torch recovery checkpoint stage IDs must be non-empty" + ) + if not isinstance(locator, TorchCheckpointLocator): + raise AlgorithmConfigurationError( + "Torch recovery stage checkpoints must use locators" + ) + checkpoints[stage_id] = locator + if set(checkpoints) != set(completed): + raise AlgorithmConfigurationError( + "Torch recovery stage checkpoints must match completed stages" + ) + if self.active_stage_id is None: + if self.active_checkpoint is not None: + raise AlgorithmConfigurationError( + "Torch recovery active checkpoint requires active_stage_id" + ) + else: + if not isinstance(self.active_stage_id, str) or not self.active_stage_id: + raise AlgorithmConfigurationError( + "Torch recovery active_stage_id must be non-empty" + ) + if self.active_stage_id in completed: + raise AlgorithmConfigurationError( + "Torch recovery active Stage cannot be completed" + ) + if not isinstance(self.active_checkpoint, TorchCheckpointLocator): + raise AlgorithmConfigurationError( + "Torch recovery active Stage requires a checkpoint locator" + ) + object.__setattr__(self, "completed_stage_ids", completed) + object.__setattr__( + self, + "stage_checkpoints", + MappingProxyType(dict(sorted(checkpoints.items()))), + ) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "completed_stage_ids": list(self.completed_stage_ids), + "stage_checkpoints": { + stage_id: locator.to_dict() + for stage_id, locator in self.stage_checkpoints.items() + }, + "active_stage_id": self.active_stage_id, + "active_checkpoint": ( + self.active_checkpoint.to_dict() + if self.active_checkpoint is not None + else None + ), + } + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "TorchRecoveryEnvelope": + raw_checkpoints = value.get("stage_checkpoints", {}) + if not isinstance(raw_checkpoints, Mapping): + raise AlgorithmConfigurationError( + "Torch recovery stage_checkpoints must be a mapping" + ) + raw_active = value.get("active_checkpoint") + parsed_checkpoints: dict[str, TorchCheckpointLocator] = {} + for stage_id, locator in raw_checkpoints.items(): + if not isinstance(stage_id, str) or not isinstance(locator, Mapping): + raise AlgorithmConfigurationError( + "Torch recovery stage_checkpoints entries are malformed" + ) + parsed_checkpoints[stage_id] = TorchCheckpointLocator.from_dict(locator) + try: + return cls( + schema_version=value.get("schema_version", 1), + completed_stage_ids=tuple(value.get("completed_stage_ids", ())), + stage_checkpoints=parsed_checkpoints, + active_stage_id=cast(str | None, value.get("active_stage_id")), + active_checkpoint=( + TorchCheckpointLocator.from_dict(raw_active) + if isinstance(raw_active, Mapping) + else None + ), + ) + except (KeyError, TypeError, ValueError) as exc: + raise AlgorithmConfigurationError( + "Torch recovery envelope is malformed" + ) from exc + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchRankProgressStatistics: + """Typed per-rank prefix statistics restored with a Torch checkpoint.""" + + rows_processed: int = 0 + coverage_totals: Mapping[str, int] = field(default_factory=dict) + loss_numerator_total: float = 0.0 + loss_normalizer_total: float = 0.0 + metric_totals: Mapping[str, tuple[float, float]] = field(default_factory=dict) + evaluation_totals: Mapping[str, tuple[float, float]] = field(default_factory=dict) + reducer_observation: Mapping[str, object] = field(default_factory=dict) + + def __post_init__(self) -> None: + if ( + not isinstance(self.rows_processed, int) + or isinstance(self.rows_processed, bool) + or self.rows_processed < 0 + ): + raise AlgorithmConfigurationError("Torch rank rows_processed is invalid") + coverage = dict(self.coverage_totals) + if any( + not isinstance(name, str) + or not name + or not isinstance(value, int) + or isinstance(value, bool) + or value < 0 + for name, value in coverage.items() + ): + raise AlgorithmConfigurationError("Torch rank coverage totals are invalid") + + def normalize_totals( + values: Mapping[str, tuple[float, float]], + ) -> dict[str, tuple[float, float]]: + normalized: dict[str, tuple[float, float]] = {} + for name, pair in values.items(): + if ( + not isinstance(name, str) + or not name + or not isinstance(pair, (list, tuple)) + or len(pair) != 2 + ): + raise AlgorithmConfigurationError( + "Torch rank metric totals are invalid" + ) + numerator = _finite_number(pair[0], f"rank metric[{name}] numerator") + normalizer = _finite_number(pair[1], f"rank metric[{name}] normalizer") + if normalizer < 0: + raise AlgorithmConfigurationError( + "Torch rank metric normalizer is invalid" + ) + normalized[name] = (numerator, normalizer) + return normalized + + loss_numerator = _finite_number( + self.loss_numerator_total, "rank loss numerator total" + ) + loss_normalizer = _finite_number( + self.loss_normalizer_total, "rank loss normalizer total" + ) + if loss_normalizer < 0: + raise AlgorithmConfigurationError("Torch rank loss normalizer is invalid") + metric_totals = normalize_totals(self.metric_totals) + evaluation_totals = normalize_totals(self.evaluation_totals) + _validate_bounded_evidence( + self.reducer_observation, path="rank.reducer_observation" + ) + object.__setattr__( + self, "coverage_totals", MappingProxyType(dict(sorted(coverage.items()))) + ) + object.__setattr__(self, "loss_numerator_total", loss_numerator) + object.__setattr__(self, "loss_normalizer_total", loss_normalizer) + object.__setattr__( + self, "metric_totals", MappingProxyType(dict(sorted(metric_totals.items()))) + ) + object.__setattr__( + self, + "evaluation_totals", + MappingProxyType(dict(sorted(evaluation_totals.items()))), + ) + object.__setattr__( + self, "reducer_observation", deep_freeze(self.reducer_observation) + ) + + def to_dict(self) -> dict[str, Any]: + return { + "rows_processed": self.rows_processed, + "coverage_totals": dict(self.coverage_totals), + "loss_numerator_total": self.loss_numerator_total, + "loss_normalizer_total": self.loss_normalizer_total, + "metric_totals": { + name: list(pair) for name, pair in self.metric_totals.items() + }, + "evaluation_totals": { + name: list(pair) for name, pair in self.evaluation_totals.items() + }, + "reducer_observation": dict(self.reducer_observation), + } + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "TorchRankProgressStatistics": + try: + return cls( + rows_processed=value.get("rows_processed", 0), + coverage_totals=value.get("coverage_totals", {}), + loss_numerator_total=value.get("loss_numerator_total", 0.0), + loss_normalizer_total=value.get("loss_normalizer_total", 0.0), + metric_totals=value.get("metric_totals", {}), + evaluation_totals=value.get("evaluation_totals", {}), + reducer_observation=value.get("reducer_observation", {}), + ) + except (TypeError, ValueError) as exc: + raise AlgorithmConfigurationError( + "Torch rank progress statistics are malformed" + ) from exc + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchCheckpointProgress: + """Deterministic cursor state required for exact Torch recovery.""" + + epoch: int + micro_batch_cursor: int + optimizer_step: int + scheduler_step: int + accumulation_steps: int + dataset_cursor_by_rank: Mapping[str, int] = field(default_factory=dict) + shuffle_seed: int | None = None + rows_processed: int = 0 + coverage_totals: Mapping[str, int] = field(default_factory=dict) + loss_numerator_total: float = 0.0 + loss_normalizer_total: float = 0.0 + metric_totals: Mapping[str, tuple[float, float]] = field(default_factory=dict) + evaluation_totals: Mapping[str, tuple[float, float]] = field(default_factory=dict) + rank_statistics: Mapping[str, TorchRankProgressStatistics] = field( + default_factory=dict + ) + epoch_scheduler_applied: bool = False + schema_version: int = 1 + + def __post_init__(self) -> None: + if self.schema_version != 1: + raise AlgorithmConfigurationError( + "unsupported Torch checkpoint progress version" + ) + for name in ( + "epoch", + "micro_batch_cursor", + "optimizer_step", + "scheduler_step", + "accumulation_steps", + ): + value = getattr(self, name) + if ( + not isinstance(value, int) + or isinstance(value, bool) + or value < 0 + or (name == "accumulation_steps" and value < 1) + ): + raise AlgorithmConfigurationError( + f"Torch checkpoint progress {name} is invalid" + ) + cursors: dict[str, int] = {} + for rank, cursor in self.dataset_cursor_by_rank.items(): + if ( + not isinstance(rank, str) + or not rank + or not isinstance(cursor, int) + or isinstance(cursor, bool) + or cursor < 0 + ): + raise AlgorithmConfigurationError( + "Torch checkpoint dataset cursors are invalid" + ) + cursors[rank] = cursor + if self.shuffle_seed is not None and ( + not isinstance(self.shuffle_seed, int) + or isinstance(self.shuffle_seed, bool) + ): + raise AlgorithmConfigurationError( + "Torch checkpoint shuffle_seed is invalid" + ) + if ( + not isinstance(self.rows_processed, int) + or isinstance(self.rows_processed, bool) + or self.rows_processed < 0 + ): + raise AlgorithmConfigurationError( + "Torch checkpoint rows_processed is invalid" + ) + coverage = dict(self.coverage_totals) + if any( + not isinstance(name, str) + or not name + or not isinstance(value, int) + or isinstance(value, bool) + or value < 0 + for name, value in coverage.items() + ): + raise AlgorithmConfigurationError( + "Torch checkpoint coverage totals are invalid" + ) + loss_numerator = _finite_number( + self.loss_numerator_total, "checkpoint loss numerator total" + ) + loss_normalizer = _finite_number( + self.loss_normalizer_total, "checkpoint loss normalizer total" + ) + if loss_normalizer < 0: + raise AlgorithmConfigurationError( + "Torch checkpoint loss normalizer is invalid" + ) + + def normalize_totals( + values: Mapping[str, tuple[float, float]], + ) -> dict[str, tuple[float, float]]: + normalized: dict[str, tuple[float, float]] = {} + for name, pair in values.items(): + if ( + not isinstance(name, str) + or not name + or not isinstance(pair, (list, tuple)) + or len(pair) != 2 + ): + raise AlgorithmConfigurationError( + "Torch checkpoint metric totals are invalid" + ) + numerator = _finite_number( + pair[0], f"checkpoint metric[{name}] numerator" + ) + normalizer = _finite_number( + pair[1], f"checkpoint metric[{name}] normalizer" + ) + if normalizer < 0: + raise AlgorithmConfigurationError( + "Torch checkpoint metric normalizer is invalid" + ) + normalized[name] = (numerator, normalizer) + return normalized + + metric_totals = normalize_totals(self.metric_totals) + evaluation_totals = normalize_totals(self.evaluation_totals) + rank_statistics = dict(self.rank_statistics) + if len(rank_statistics) > 1024 or any( + not isinstance(rank, str) + or not rank + or not isinstance(stats, TorchRankProgressStatistics) + for rank, stats in rank_statistics.items() + ): + raise AlgorithmConfigurationError( + "Torch checkpoint rank statistics are invalid" + ) + if not isinstance(self.epoch_scheduler_applied, bool): + raise AlgorithmConfigurationError( + "Torch checkpoint epoch_scheduler_applied must be boolean" + ) + object.__setattr__( + self, + "dataset_cursor_by_rank", + MappingProxyType(dict(sorted(cursors.items()))), + ) + object.__setattr__(self, "rows_processed", self.rows_processed) + object.__setattr__( + self, "coverage_totals", MappingProxyType(dict(sorted(coverage.items()))) + ) + object.__setattr__(self, "loss_numerator_total", loss_numerator) + object.__setattr__(self, "loss_normalizer_total", loss_normalizer) + object.__setattr__( + self, "metric_totals", MappingProxyType(dict(sorted(metric_totals.items()))) + ) + object.__setattr__( + self, + "evaluation_totals", + MappingProxyType(dict(sorted(evaluation_totals.items()))), + ) + object.__setattr__( + self, + "rank_statistics", + MappingProxyType(dict(sorted(rank_statistics.items()))), + ) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "epoch": self.epoch, + "micro_batch_cursor": self.micro_batch_cursor, + "optimizer_step": self.optimizer_step, + "scheduler_step": self.scheduler_step, + "accumulation_steps": self.accumulation_steps, + "dataset_cursor_by_rank": dict(self.dataset_cursor_by_rank), + "shuffle_seed": self.shuffle_seed, + "rows_processed": self.rows_processed, + "coverage_totals": dict(self.coverage_totals), + "loss_numerator_total": self.loss_numerator_total, + "loss_normalizer_total": self.loss_normalizer_total, + "metric_totals": { + name: list(pair) for name, pair in self.metric_totals.items() + }, + "evaluation_totals": { + name: list(pair) for name, pair in self.evaluation_totals.items() + }, + "rank_statistics": { + rank: statistics.to_dict() + for rank, statistics in self.rank_statistics.items() + }, + "epoch_scheduler_applied": self.epoch_scheduler_applied, + } + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "TorchCheckpointProgress": + raw_cursors = value.get("dataset_cursor_by_rank", {}) + if not isinstance(raw_cursors, Mapping): + raise AlgorithmConfigurationError( + "Torch checkpoint dataset cursors must be a mapping" + ) + raw_statistics = value.get("rank_statistics", {}) + if not isinstance(raw_statistics, Mapping) or any( + not isinstance(rank, str) or not isinstance(statistics, Mapping) + for rank, statistics in raw_statistics.items() + ): + raise AlgorithmConfigurationError( + "Torch checkpoint rank statistics must be a typed mapping" + ) + try: + return cls( + schema_version=value.get("schema_version", 1), + epoch=value["epoch"], + micro_batch_cursor=value["micro_batch_cursor"], + optimizer_step=value["optimizer_step"], + scheduler_step=value["scheduler_step"], + accumulation_steps=value["accumulation_steps"], + dataset_cursor_by_rank=dict(raw_cursors), + shuffle_seed=value.get("shuffle_seed"), + rows_processed=value.get("rows_processed", 0), + coverage_totals=value.get("coverage_totals", {}), + loss_numerator_total=value.get("loss_numerator_total", 0.0), + loss_normalizer_total=value.get("loss_normalizer_total", 0.0), + metric_totals=value.get("metric_totals", {}), + evaluation_totals=value.get("evaluation_totals", {}), + rank_statistics={ + rank: TorchRankProgressStatistics.from_dict(statistics) + for rank, statistics in raw_statistics.items() + }, + epoch_scheduler_applied=value.get("epoch_scheduler_applied", False), + ) + except (KeyError, TypeError, ValueError) as exc: + raise AlgorithmConfigurationError( + "Torch checkpoint progress is malformed" + ) from exc + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchCheckpointDescriptor: + """Validated manifest embedded in every portable Torch checkpoint.""" + + schema_version: int + identity: TorchStageRunIdentity + run_config_name: str + state_layout: str + world_size: int + completed_step: int + policy_digest: str + execution_plan_digest: str + input_binding_digest: str + implementation_code_digest: str + payload_files: Mapping[str, str] + adapter_identity: str | None = None + resume_supported: bool = True + same_world_size_resume: bool | None = True + torch_runtime_api_version: int = 1 + + def __post_init__(self) -> None: + if self.schema_version != 1: + raise AlgorithmConfigurationError("unsupported Torch checkpoint descriptor") + if not isinstance(self.identity, TorchStageRunIdentity): + raise AlgorithmConfigurationError("checkpoint identity is invalid") + if ( + self.torch_runtime_api_version != 1 + or self.identity.torch_runtime_api_version != 1 + ): + raise AlgorithmConfigurationError( + "Torch checkpoint runtime API version must be exactly 1" + ) + for field_name in ("run_config_name", "input_binding_digest"): + value = getattr(self, field_name) + if not isinstance(value, str) or not value: + raise AlgorithmConfigurationError(f"{field_name} must be non-empty") + if self.run_config_name != self.identity.run_config_name: + raise AlgorithmConfigurationError( + "checkpoint run_config_name does not match Stage identity" + ) + if self.state_layout not in {"replicated", "component", "sharded"}: + raise AlgorithmConfigurationError("checkpoint state_layout is invalid") + for field_name in ("world_size", "completed_step"): + value = getattr(self, field_name) + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise AlgorithmConfigurationError(f"{field_name} must be non-negative") + if self.world_size < 1: + raise AlgorithmConfigurationError("world_size must be positive") + for field_name in ( + "policy_digest", + "execution_plan_digest", + "input_binding_digest", + "implementation_code_digest", + ): + _digest_value(getattr(self, field_name), field_name) + if self.identity.policy_digest != self.policy_digest: + raise AlgorithmConfigurationError( + "checkpoint Policy digest does not match identity" + ) + if self.identity.execution_plan_digest != self.execution_plan_digest: + raise AlgorithmConfigurationError( + "checkpoint execution plan digest does not match identity" + ) + if self.identity.implementation_code_digest != self.implementation_code_digest: + raise AlgorithmConfigurationError( + "checkpoint implementation code digest does not match identity" + ) + files = dict(self.payload_files) + if not files: + raise AlgorithmConfigurationError("checkpoint payload_files are required") + for name, digest in files.items(): + if ( + not isinstance(name, str) + or not name + or name.startswith("/") + or "\\" in name + or ".." in name.split("/") + ): + raise AlgorithmConfigurationError( + "checkpoint payload file path is unsafe" + ) + _digest_value(digest, f"payload_files[{name}]") + if not isinstance(self.resume_supported, bool): + raise AlgorithmConfigurationError("resume_supported must be boolean") + if self.same_world_size_resume is not None and not isinstance( + self.same_world_size_resume, bool + ): + raise AlgorithmConfigurationError("same_world_size_resume must be boolean") + if self.resume_supported and self.same_world_size_resume is not True: + raise AlgorithmConfigurationError( + "supported recovery requires same world size" + ) + if not self.resume_supported and self.same_world_size_resume is not None: + raise AlgorithmConfigurationError( + "unsupported recovery must omit same_world_size_resume" + ) + object.__setattr__( + self, "payload_files", MappingProxyType(dict(sorted(files.items()))) + ) + + @property + def digest(self) -> str: + return _digest(self.to_dict()) + + def to_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "schema_version": self.schema_version, + "identity": self.identity.to_dict(), + "run_config_name": self.run_config_name, + "state_layout": self.state_layout, + "world_size": self.world_size, + "completed_step": self.completed_step, + "policy_digest": self.policy_digest, + "execution_plan_digest": self.execution_plan_digest, + "input_binding_digest": self.input_binding_digest, + "implementation_code_digest": self.implementation_code_digest, + "payload_files": dict(self.payload_files), + "adapter_identity": self.adapter_identity, + "resume_supported": self.resume_supported, + "torch_runtime_api_version": self.torch_runtime_api_version, + } + if self.same_world_size_resume is not None: + payload["same_world_size_resume"] = self.same_world_size_resume + return payload + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "TorchCheckpointDescriptor": + try: + identity = TorchStageRunIdentity.from_dict(value["identity"]) + resume_supported = value.get("resume_supported", True) + return cls( + schema_version=value["schema_version"], + identity=identity, + run_config_name=value["run_config_name"], + state_layout=value["state_layout"], + world_size=value["world_size"], + completed_step=value["completed_step"], + policy_digest=value["policy_digest"], + execution_plan_digest=value["execution_plan_digest"], + input_binding_digest=value["input_binding_digest"], + implementation_code_digest=value["implementation_code_digest"], + payload_files=value["payload_files"], + adapter_identity=value.get("adapter_identity"), + resume_supported=resume_supported, + same_world_size_resume=value.get( + "same_world_size_resume", True if resume_supported else None + ), + torch_runtime_api_version=value.get("torch_runtime_api_version", 1), + ) + except KeyError as exc: + raise AlgorithmConfigurationError( + f"Torch checkpoint descriptor is missing {exc.args[0]!r}" + ) from exc + + +__all__ = [ + "TorchAccumulationWindow", + "TorchBackwardContext", + "TorchBackwardResult", + "TorchCheckpointPayloadDraft", + "TorchCheckpointDescriptor", + "TorchCheckpointLocator", + "TorchRecoveryEnvelope", + "TorchCheckpointProgress", + "TorchRankProgressStatistics", + "TorchCheckpointRef", + "TorchCompositeGlobalState", + "TorchCompositeLossContribution", + "TorchGlobalLossContext", + "TorchGlobalLossReducer", + "TorchGlobalLossReduction", + "TorchLossContribution", + "TorchMetricContribution", + "TorchMetricPolicy", + "TorchMetricReductionContext", + "TorchMetricReductionResult", + "TorchPreflightLease", + "TorchPreflightTokenData", + "TorchStageRunIdentity", + "TorchStepLoss", + "TorchRuntimeExecutionEnvelope", + "TorchWorkerControlEnvelope", + "apply_torch_loss_backward", + "claim_torch_run_directory", + "describe_torch_checkpoint", + "invoke_torch_global_loss_reducer", + "reduce_torch_metrics", + "report_torch_checkpoint", + "torch_run_config_name", + "validate_torch_retry_identity", +] diff --git a/src/tributo/algorithms/composition.py b/src/tributo/algorithms/composition.py index 010fc65..0ef153e 100644 --- a/src/tributo/algorithms/composition.py +++ b/src/tributo/algorithms/composition.py @@ -28,7 +28,6 @@ def build_algorithm_dispatcher( ) from tributo.integrations.algorithm_runtimes.collective import ( RayTrainCollectiveRuntime, - RayTrainRecipeV2Runtime, ) from tributo.integrations.algorithm_runtimes.framework_native import ( FrameworkNativeRuntime, @@ -46,6 +45,9 @@ def build_algorithm_dispatcher( RayParallelUnitRuntime, ) from tributo.integrations.algorithm_runtimes.ray_task import RayTaskRuntime + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + RayTrainTorchRuntime, + ) from tributo.training.registry import get_execution_registry resolver = IngestionInputResolver( @@ -53,10 +55,10 @@ def build_algorithm_dispatcher( accepted_handle_kinds=("ray_data",), ) input_adapter = IngestionInputRuntimeAdapter() - runtimes = ( + runtimes: tuple[Any, ...] = ( RayTaskRuntime(), RayTrainCollectiveRuntime(), - RayTrainRecipeV2Runtime(), + RayTrainTorchRuntime(), RayMapReduceRuntime(), FrameworkNativeRuntime(), RayJoblibEstimatorRuntime(), diff --git a/src/tributo/algorithms/core/builder.py b/src/tributo/algorithms/core/builder.py index 8a5445e..3da417b 100644 --- a/src/tributo/algorithms/core/builder.py +++ b/src/tributo/algorithms/core/builder.py @@ -30,6 +30,10 @@ ResultPolicy, RuntimeBinding, RuntimeTopology, + SingleStageTorchPlan, + TorchDatasetRoute, + TorchPolicy, + TorchStageSpec, WorkerRange, WorkerResources, ) @@ -112,6 +116,7 @@ def from_distributed_algorithm( | JoblibEstimatorPolicy | ParallelEnsemblePolicy | IterativeOptimizationPolicy + | TorchPolicy ), package_name: str, package_version: str, @@ -240,7 +245,7 @@ def from_distributed_algorithm( ) @staticmethod - def from_torch_recipe( + def from_torch( *, spec: AlgorithmSpec, implementation_id: str, @@ -254,6 +259,7 @@ def from_torch_recipe( package_name: str, package_version: str, tributo_version_spec: str, + policy: TorchPolicy | None = None, backend: Literal["auto", "gloo", "nccl"] = "auto", distributed_min_workers: int = 2, stability: Literal["alpha", "beta", "stable"] = "alpha", @@ -264,45 +270,72 @@ def from_torch_recipe( is_default: bool = False, code_digest: str | None = None, contract_bindings: ContractBindingSet | None = None, - descriptor_api_version: int = 1, + descriptor_api_version: int = 2, ) -> DistributedAlgorithmDescriptor: - """Lower four PyTorch factories to the existing Ray collective runtime. - - The referenced class must subclass ``TorchTrainingRecipe`` and have a - no-argument constructor. The ordinary recipe configuration surface is - deliberately fixed to model, loss, optimizer, metrics, training, ray, - and output namespaces so algorithm code cannot smuggle deployment - settings into the worker loop. - """ - try: - normalized_reducers = { - name: MetricReduction(reduction) - for name, reduction in metric_reducers.items() - } - except (TypeError, ValueError) as exc: + """Build a registration for the unified Core-owned Torch Runtime.""" + if code_digest is None: raise AlgorithmConfigurationError( - "Torch recipe metric reducer is invalid" - ) from exc + "Torch implementations require an explicit code_digest" + ) + normalized_reducers = { + name: MetricReduction(reduction) + for name, reduction in metric_reducers.items() + } if "train_loss" in normalized_reducers and ( normalized_reducers["train_loss"] is not MetricReduction.SUM_COUNT ): raise AlgorithmConfigurationError( - "Torch recipe train_loss uses the fixed sum_count reducer" + "Torch train_loss uses the fixed sum_count reducer" ) normalized_reducers["train_loss"] = MetricReduction.SUM_COUNT + resolved_policy = policy or TorchPolicy( + torch_runtime_api_version=1, + loop_owner="core_recipe", + parallelism_id="torch.ddp.replicated", + dataset_routing=( + TorchDatasetRoute( + role="train", + mode="split_exact", + required=True, + min_total_rows_if_present=1, + min_rows_per_worker=1, + empty_rank_policy="reject", + ), + ), + execution_plan=SingleStageTorchPlan( + stage=TorchStageSpec( + stage_id="train", + worker_loop_ref=( + "tributo.integrations.algorithm_runtimes.ray_train_torch:" + "torch_recipe_train_loop_per_worker" + ), + input_roles=("train",), + ) + ), + state_layout="replicated", + metric_reducers=normalized_reducers, + backend=backend, + resume_supported=True, + same_world_size_resume=True, + ) + if resolved_policy.backend != backend and policy is not None: + raise AlgorithmConfigurationError( + "Torch Policy backend conflicts with the requested backend" + ) return AlgorithmBuilder.from_distributed_algorithm( spec=spec, implementation_id=implementation_id, implementation_version=implementation_version, implementation=recipe, executable_factory=( - "tributo.integrations.algorithm_runtimes.torch_recipe:" - "create_torch_recipe_algorithm" + "tributo.integrations.algorithm_runtimes.ray_train_torch:" + "create_torch_algorithm" ), distribution=package_name, framework="pytorch", environment=environment, allowed_config_keys=( + "data", "loss", "metrics", "model", @@ -311,24 +344,18 @@ def from_torch_recipe( "ray", "training", ), - strategy=DistributionStrategy.RAY_TRAIN_COLLECTIVE, + strategy=DistributionStrategy.RAY_TRAIN_TORCH, supported_worker_range=supported_worker_range, supported_execution_profiles=supported_execution_profiles, resources_per_worker=resources_per_worker, - policy=CollectivePolicy( - backend=backend, - metric_reducers=normalized_reducers, - checkpoint_owner_rank=0, - same_world_size_resume=True, - rank_seeded=True, - ), + policy=resolved_policy, package_name=package_name, package_version=package_version, tributo_version_spec=tributo_version_spec, result_policy=ResultPolicy.BUNDLE_REQUIRED, exporter=( - "tributo.integrations.algorithm_runtimes.torch_recipe:" - "export_torch_recipe_result" + "tributo.integrations.algorithm_runtimes.ray_train_torch:" + "export_ray_train_torch_result" ), flavor_id="onnx-runtime-v1", distributed_min_workers=distributed_min_workers, @@ -344,140 +371,117 @@ def from_torch_recipe( ) @staticmethod - def from_joblib_estimator_recipe( + def from_torch_adapter( *, spec: AlgorithmSpec, implementation_id: str, implementation_version: str, - recipe: str, + adapter: str, environment: EnvironmentSpec, - allowed_config_keys: tuple[str, ...], + metric_reducers: Mapping[str, MetricReduction], supported_worker_range: WorkerRange, supported_execution_profiles: tuple[ExecutionProfile, ...], resources_per_worker: WorkerResources, package_name: str, package_version: str, tributo_version_spec: str, - policy: JoblibEstimatorPolicy | None = None, - input_compatibility: BackendInputCompatibility | None = None, - result_policy: ResultPolicy = ResultPolicy.BUNDLE_REQUIRED, - exporter: str | None = None, - flavor_id: str | None = None, + policy: TorchPolicy, distributed_min_workers: int = 2, + stability: Literal["alpha", "beta", "stable"] = "alpha", + tested: bool = False, + supported: bool = False, + validated_execution_profiles: tuple[ExecutionProfile, ...] = (), + limitations: tuple[str, ...] = (), + is_default: bool = False, + code_digest: str | None = None, contract_bindings: ContractBindingSet | None = None, descriptor_api_version: int = 2, - is_default: bool = False, ) -> DistributedAlgorithmDescriptor: - """Lower estimator mathematics to the Core Ray Joblib Runtime.""" - return AlgorithmBuilder.from_distributed_algorithm( + """Build a registration for a framework-owned ``RayTorchAdapter``.""" + if policy.loop_owner != "adapter": + raise AlgorithmConfigurationError( + "Torch adapter registrations require policy.loop_owner='adapter'" + ) + if not policy.resume_supported and policy.same_world_size_resume is not None: + raise AlgorithmConfigurationError( + "adapter registrations must omit same_world_size_resume when recovery is disabled" + ) + return AlgorithmBuilder.from_torch( spec=spec, implementation_id=implementation_id, implementation_version=implementation_version, - implementation=recipe, - executable_factory=( - "tributo.integrations.algorithm_runtimes.decomposition:create_algorithm" - ), - distribution=package_name, - framework="sklearn", + recipe=adapter, environment=environment, - allowed_config_keys=allowed_config_keys, - strategy=DistributionStrategy.RAY_JOBLIB_ESTIMATOR, + metric_reducers=metric_reducers, supported_worker_range=supported_worker_range, supported_execution_profiles=supported_execution_profiles, resources_per_worker=resources_per_worker, - policy=policy or JoblibEstimatorPolicy(), package_name=package_name, package_version=package_version, tributo_version_spec=tributo_version_spec, - result_policy=result_policy, - input_compatibility=input_compatibility, - exporter=exporter, - flavor_id=flavor_id, + policy=policy, distributed_min_workers=distributed_min_workers, + stability=stability, + tested=tested, + supported=supported, + validated_execution_profiles=validated_execution_profiles, + limitations=limitations, + is_default=is_default, + code_digest=code_digest, contract_bindings=contract_bindings, descriptor_api_version=descriptor_api_version, - is_default=is_default, ) @staticmethod - def from_training_recipe_v2( + def from_joblib_estimator_recipe( *, spec: AlgorithmSpec, implementation_id: str, implementation_version: str, recipe: str, environment: EnvironmentSpec, - metric_reducers: Mapping[str, MetricReduction], + allowed_config_keys: tuple[str, ...], supported_worker_range: WorkerRange, supported_execution_profiles: tuple[ExecutionProfile, ...], resources_per_worker: WorkerResources, package_name: str, package_version: str, tributo_version_spec: str, - contract_bindings: ContractBindingSet, - backend: Literal["auto", "gloo", "nccl"] = "auto", + policy: JoblibEstimatorPolicy | None = None, + input_compatibility: BackendInputCompatibility | None = None, + result_policy: ResultPolicy = ResultPolicy.BUNDLE_REQUIRED, + exporter: str | None = None, + flavor_id: str | None = None, distributed_min_workers: int = 2, + contract_bindings: ContractBindingSet | None = None, descriptor_api_version: int = 2, is_default: bool = False, ) -> DistributedAlgorithmDescriptor: - """Lower TrainingRecipeV2 Step/Plan Hooks to the Core DDP loop.""" - try: - normalized_reducers = { - name: MetricReduction(reduction) - for name, reduction in metric_reducers.items() - } - except (TypeError, ValueError) as exc: - raise AlgorithmConfigurationError( - "TrainingRecipeV2 metric reducer is invalid" - ) from exc - if "train_loss" in normalized_reducers and ( - normalized_reducers["train_loss"] is not MetricReduction.SUM_COUNT - ): - raise AlgorithmConfigurationError( - "TrainingRecipeV2 train_loss uses the fixed sum_count reducer" - ) - normalized_reducers["train_loss"] = MetricReduction.SUM_COUNT + """Lower estimator mathematics to the Core Ray Joblib Runtime.""" return AlgorithmBuilder.from_distributed_algorithm( spec=spec, implementation_id=implementation_id, implementation_version=implementation_version, implementation=recipe, executable_factory=( - "tributo.integrations.algorithm_runtimes.torch_recipe:" - "create_torch_recipe_algorithm" + "tributo.integrations.algorithm_runtimes.decomposition:create_algorithm" ), distribution=package_name, - framework="pytorch", + framework="sklearn", environment=environment, - allowed_config_keys=( - "loss", - "metrics", - "model", - "optimizer", - "output", - "ray", - "training", - ), - strategy=DistributionStrategy.RAY_TRAIN_RECIPE_V2, + allowed_config_keys=allowed_config_keys, + strategy=DistributionStrategy.RAY_JOBLIB_ESTIMATOR, supported_worker_range=supported_worker_range, supported_execution_profiles=supported_execution_profiles, resources_per_worker=resources_per_worker, - policy=CollectivePolicy( - backend=backend, - metric_reducers=normalized_reducers, - checkpoint_owner_rank=0, - same_world_size_resume=True, - rank_seeded=True, - ), + policy=policy or JoblibEstimatorPolicy(), package_name=package_name, package_version=package_version, tributo_version_spec=tributo_version_spec, - result_policy=ResultPolicy.BUNDLE_REQUIRED, - exporter=( - "tributo.integrations.algorithm_runtimes.torch_recipe:" - "export_torch_recipe_result" - ), - flavor_id="onnx-runtime-v1", + result_policy=result_policy, + input_compatibility=input_compatibility, + exporter=exporter, + flavor_id=flavor_id, distributed_min_workers=distributed_min_workers, contract_bindings=contract_bindings, descriptor_api_version=descriptor_api_version, diff --git a/src/tributo/algorithms/core/dispatcher.py b/src/tributo/algorithms/core/dispatcher.py index b194508..636f452 100644 --- a/src/tributo/algorithms/core/dispatcher.py +++ b/src/tributo/algorithms/core/dispatcher.py @@ -6,7 +6,7 @@ import uuid from collections.abc import Mapping from dataclasses import replace -from typing import cast +from typing import Any, cast from tributo._common.immutable import deep_thaw from tributo.algorithms.api import ( @@ -16,10 +16,14 @@ AlgorithmRequest, AlgorithmRunResult, ArtifactDraft, + DistributionStrategy, ExecutionReceipt, ExecutionRequest, ResolvedAlgorithmPlan, StateCoordinationEvidence, + TorchExecutionEvidence, + TorchPreflightLease, + TorchRuntimeExecutionEnvelope, WorkerExecutionEvidence, WorkerExecutionResult, WorkerResources, @@ -36,6 +40,7 @@ ResolvedInputLease, RuntimeExecutionEnvelope, RuntimeInputBinding, + TorchRuntimePreflight, WorkerInputPayload, WorkerInputPayloadSet, ) @@ -65,10 +70,20 @@ def execute( context: InputExecutionContext, artifacts: tuple[ArtifactDraft, ...] = (), cancelled: bool = False, + run_id: str | None = None, + torch_preflight_lease: TorchPreflightLease | None = None, ) -> AlgorithmRunResult: """Open input, invoke the selected Runtime, and close in reverse order.""" plan.validate_integrity() - run_id = uuid.uuid4().hex + run_id = run_id or uuid.uuid4().hex + is_torch = ( + plan.distribution_spec is not None + and plan.distribution_spec.strategy is DistributionStrategy.RAY_TRAIN_TORCH + ) + if is_torch and torch_preflight_lease is None: + raise AlgorithmConfigurationError( + "Torch execution requires a preflight lease" + ) try: runtime = self._runtimes[plan.runtime.runtime_id] except KeyError as exc: @@ -100,14 +115,20 @@ def execute( leases.append((binding.name, lease)) role_bindings.append(input_adapter.bind(lease, plan)) runtime_binding = self._combine_role_bindings(plan, role_bindings) - runtime_result = runtime.execute( - RuntimeExecutionEnvelope( - run_id=run_id, - plan=plan, - input_payloads=runtime_binding.payloads, - artifacts=artifacts, - cancelled=cancelled, + base_envelope = RuntimeExecutionEnvelope( + run_id=run_id, + plan=plan, + input_payloads=runtime_binding.payloads, + artifacts=artifacts, + cancelled=cancelled, + ) + runtime_result = cast(Any, runtime).execute( + TorchRuntimeExecutionEnvelope( + base=base_envelope, + preflight_lease=cast(TorchPreflightLease, torch_preflight_lease), ) + if is_torch + else base_envelope ) if not isinstance(runtime_result, WorkerExecutionResult) or not isinstance( runtime_result.execution, AlgorithmExecutionResult @@ -162,6 +183,12 @@ def execute( lease.close() except Exception as exc: cleanup_errors.append(exc) + if is_torch and torch_preflight_lease is not None: + try: + if torch_preflight_lease.state != "consumed": + torch_preflight_lease.close() + except Exception as exc: + cleanup_errors.append(exc) if primary_error is not None: for cleanup_error in cleanup_errors: @@ -216,6 +243,48 @@ def execute( execution_receipt=execution_receipt, ) + def torch_preflight( + self, + plan: ResolvedAlgorithmPlan, + *, + run_id: str, + invocation_id: str, + ) -> TorchPreflightLease: + """Run the Torch-only environment check before any input or Ray lease.""" + if ( + plan.distribution_spec is None + or plan.distribution_spec.strategy + is not DistributionStrategy.RAY_TRAIN_TORCH + ): + raise AlgorithmConfigurationError( + "Torch preflight requires RAY_TRAIN_TORCH" + ) + try: + runtime = self._runtimes[plan.runtime.runtime_id] + except KeyError as exc: + raise AlgorithmConfigurationError( + f"missing execution component for {exc.args[0]!r}" + ) from exc + if not isinstance(runtime, TorchRuntimePreflight): + raise AlgorithmConfigurationError("selected Torch runtime has no preflight") + return runtime.preflight(plan, run_id, invocation_id) + + @staticmethod + def claim_torch_preflight( + plan: ResolvedAlgorithmPlan, + *, + run_id: str, + invocation_id: str, + lease: TorchPreflightLease, + ) -> None: + """Claim a preflight lease before opening Runtime or inputs.""" + lease.claim( + run_id=run_id, + invocation_id=invocation_id, + plan_digest=plan.plan_id, + runtime_id=plan.runtime.runtime_id, + ) + @staticmethod def _combine_role_bindings( plan: ResolvedAlgorithmPlan, @@ -286,6 +355,111 @@ def _execution_receipt( "worker evidence entries must be mappings" ) state = StateCoordinationEvidence.from_dict(state_value) + torch_evidence = None + if plan.distribution_spec.strategy is DistributionStrategy.RAY_TRAIN_TORCH: + raw_torch_evidence = metadata.get("torch_evidence") + if not isinstance(raw_torch_evidence, Mapping): + raise AlgorithmConfigurationError( + "Torch runtime did not return TorchExecutionEvidence" + ) + torch_evidence = TorchExecutionEvidence.from_dict(raw_torch_evidence) + if torch_evidence.identity.run_id != run_id: + raise AlgorithmConfigurationError( + "Torch execution evidence run identity does not match invocation" + ) + policy = cast(Any, plan.distribution_spec.policy) + if ( + not hasattr(policy, "digest") + or torch_evidence.policy_digest != policy.digest + ): + raise AlgorithmConfigurationError( + "Torch execution evidence policy digest does not match plan" + ) + if ( + torch_evidence.identity.execution_plan_digest + != policy.execution_plan.digest + ): + raise AlgorithmConfigurationError( + "Torch execution evidence execution plan digest does not match plan" + ) + if ( + torch_evidence.state_layout != policy.state_layout + or torch_evidence.parallelism_id != policy.parallelism_id + ): + raise AlgorithmConfigurationError( + "Torch execution evidence state layout/topology does not match policy" + ) + if torch_evidence.replicated_state is not None and ( + state_value.get("global_model_digest") + != torch_evidence.replicated_state.global_model_digest + ): + raise AlgorithmConfigurationError( + "Torch state and evidence global model digests differ" + ) + if torch_evidence.stages and torch_evidence.final_stage_id is not None: + final_state = next( + stage.state_digest + for stage in torch_evidence.stages + if stage.stage_id == torch_evidence.final_stage_id + ) + if state_value.get("global_model_digest") != final_state: + raise AlgorithmConfigurationError( + "Torch component state digest does not match final Stage" + ) + expected_stages = tuple( + stage.stage_id for stage in policy.execution_plan.stages + ) + observed_stages = tuple( + stage.stage_id for stage in torch_evidence.stages + ) + if torch_evidence.state_layout == "component": + if observed_stages != expected_stages: + raise AlgorithmConfigurationError( + "Torch execution evidence Stage set/order does not match plan" + ) + if ( + torch_evidence.final_stage_id + != policy.execution_plan.final_stage_id + ): + raise AlgorithmConfigurationError( + "Torch execution evidence final Stage does not match plan" + ) + final_stage = next( + stage + for stage in policy.execution_plan.stages + if stage.stage_id == policy.execution_plan.final_stage_id + ) + if {role.role for role in torch_evidence.roles} != set( + final_stage.input_roles + ): + raise AlgorithmConfigurationError( + "Torch execution evidence roles do not match final Stage" + ) + if policy.global_loss_reducer_ref is None: + if any( + value is not None + for value in ( + torch_evidence.reducer_id, + torch_evidence.reducer_api_version, + torch_evidence.reducer_schema_id, + torch_evidence.reducer_code_digest, + ) + ): + raise AlgorithmConfigurationError( + "Torch execution evidence declares an unexpected reducer" + ) + elif ( + torch_evidence.reducer_id is None + or torch_evidence.reducer_api_version + != policy.global_loss_reducer_api_version + or torch_evidence.reducer_schema_id + != policy.composite_loss_schema_id + or torch_evidence.reducer_code_digest + != policy.global_loss_reducer_code_digest + ): + raise AlgorithmConfigurationError( + "Torch execution evidence reducer does not match policy" + ) input_complete = metadata.get("input_complete") driver_rows = metadata.get("driver_materialized_training_rows") if not isinstance(input_complete, bool): @@ -327,6 +501,7 @@ def _execution_receipt( driver_materialized_training_rows=driver_rows, artifact_ids=tuple(artifact_ids), cluster_resources={}, + torch_evidence=torch_evidence, ) @@ -388,13 +563,46 @@ def execute_plan( ) -> AlgorithmRunResult: """Execute one already validated plan through the normal lifecycle.""" plan.validate_integrity() - if plan.runtime.execution_profile is None: - return self._coordinator.execute( + is_torch = ( + plan.distribution_spec is not None + and plan.distribution_spec.strategy is DistributionStrategy.RAY_TRAIN_TORCH + ) + if is_torch and cancelled: + raise AlgorithmExecutionError("Torch execution was cancelled") + run_id = uuid.uuid4().hex if is_torch else None + torch_lease: TorchPreflightLease | None = None + if is_torch: + torch_run_id = cast(str, run_id) + invocation_id = uuid.uuid4().hex + torch_lease = self._coordinator.torch_preflight( plan, - context, - artifacts, - cancelled=cancelled, + run_id=torch_run_id, + invocation_id=invocation_id, ) + try: + self._coordinator.claim_torch_preflight( + plan, + run_id=torch_run_id, + invocation_id=invocation_id, + lease=torch_lease, + ) + except BaseException: + torch_lease.close() + raise + if plan.runtime.execution_profile is None: + try: + return self._coordinator.execute( + plan, + context, + artifacts, + cancelled=cancelled, + run_id=run_id, + torch_preflight_lease=torch_lease, + ) + except BaseException: + if torch_lease is not None and torch_lease.state != "consumed": + torch_lease.close() + raise if plan.distribution_spec is None: raise AlgorithmConfigurationError( "formal execution profile requires a DistributionSpec" @@ -405,27 +613,34 @@ def execute_plan( memory_bytes=getattr(plan.runtime, "memory_bytes", None), custom=plan.runtime.custom_resources, ) - with self._runtime_manager.open( - plan.runtime.execution_profile, - resources_per_worker=resources, - worker_count=plan.runtime.worker_count, - ) as runtime_session: - result = self._coordinator.execute( - plan, - context, - artifacts, - cancelled=cancelled, - ) - receipt = result.execution_receipt - if receipt is None: - return result - updated_receipt = replace( - cast(ExecutionReceipt, receipt), - cluster_resources=dict(runtime_session.cluster_resources), - runtime_owned=runtime_session.runtime_owned, - resource_preflight=runtime_session.resource_preflight, - ) - return replace(result, execution_receipt=updated_receipt) + try: + with self._runtime_manager.open( + plan.runtime.execution_profile, + resources_per_worker=resources, + worker_count=plan.runtime.worker_count, + ) as runtime_session: + result = self._coordinator.execute( + plan, + context, + artifacts, + cancelled=cancelled, + run_id=run_id, + torch_preflight_lease=torch_lease, + ) + receipt = result.execution_receipt + if receipt is None: + return result + updated_receipt = replace( + cast(ExecutionReceipt, receipt), + cluster_resources=dict(runtime_session.cluster_resources), + runtime_owned=runtime_session.runtime_owned, + resource_preflight=runtime_session.resource_preflight, + ) + return replace(result, execution_receipt=updated_receipt) + except BaseException: + if torch_lease is not None and torch_lease.state != "consumed": + torch_lease.close() + raise __all__ = ["AlgorithmDispatcher", "AlgorithmRunCoordinator"] diff --git a/src/tributo/algorithms/core/planner.py b/src/tributo/algorithms/core/planner.py index 8c924c8..250fbea 100644 --- a/src/tributo/algorithms/core/planner.py +++ b/src/tributo/algorithms/core/planner.py @@ -3,13 +3,14 @@ from __future__ import annotations from collections.abc import Mapping -from typing import Any +from typing import Any, cast from tributo._common.immutable import deep_thaw from tributo.algorithms.api import ( AlgorithmConfigurationError, AlgorithmRequest, AlgorithmResolution, + DistributionStrategy, ExecutionRequest, InputBinding, InputBindingSet, @@ -19,6 +20,7 @@ ResolvedInputDescriptorSet, RuntimeBinding, RuntimeTopology, + TorchPolicy, WorkerResources, canonical_digest, ) @@ -99,6 +101,20 @@ def plan( available_resources=available_resources, ) binding_set = algorithm_request.input_bindings + if ( + registration.distribution_spec is not None + and registration.distribution_spec.strategy + is DistributionStrategy.RAY_TRAIN_TORCH + ): + policy = cast(TorchPolicy, registration.distribution_spec.policy) + routes = {route.role: route for route in policy.dataset_routing} + unknown_roles = sorted( + {binding.name for binding in binding_set.bindings} - set(routes) + ) + if unknown_roles: + raise AlgorithmConfigurationError( + f"Torch input binding role(s) are not declared by TorchPolicy: {unknown_roles}" + ) resolution_context = context or InputResolutionContext() descriptors: list[ResolvedInputDescriptor] = [] for binding in binding_set.bindings: @@ -229,7 +245,7 @@ def _validate_input_descriptor( if ( registration.distribution_spec is not None and registration.distribution_spec.input_distribution - is InputDistribution.SHARDED + in {InputDistribution.SHARDED, InputDistribution.ROLE_ROUTED} and "shardable" not in descriptor.input_capabilities ): raise AlgorithmConfigurationError( @@ -300,6 +316,11 @@ def _resolve_runtime( strategy=distribution.strategy, distribution_digest=distribution.digest, resume_from=request.resume_from, + torch_recovery=( + request.torch_recovery.to_dict() + if request.torch_recovery is not None + else None + ), ) @staticmethod diff --git a/src/tributo/algorithms/spi/__init__.py b/src/tributo/algorithms/spi/__init__.py index c35e0f3..5c76611 100644 --- a/src/tributo/algorithms/spi/__init__.py +++ b/src/tributo/algorithms/spi/__init__.py @@ -16,6 +16,7 @@ PortableRuntimeAdapter, Predictable, RuntimeExecutionEnvelope, + TorchRuntimePreflight, Transformable, ) from tributo.algorithms.spi.input import ( @@ -33,11 +34,22 @@ WorkerInputPayloadSet, ) from tributo.algorithms.spi.torch import ( - MetricPlan, - OptimizationPlan, - TorchTrainingRecipe, - TrainingRecipeV2, - TrainingStepResult, + RayTorchAdapter, + TorchArtifactContext, + TorchArtifactPlan, + TorchBatch, + TorchBatchContext, + TorchBuildContext, + TorchCheckpointContext, + TorchMetricPlan, + TorchModuleSet, + TorchOptimizationPlan, + TorchRecipe, + TorchRuntimeContext, + TorchStageContext, + TorchStepContext, + TorchStepResult, + TorchWorkerCheckpointContext, ) __all__ = [ @@ -57,20 +69,32 @@ "InputRuntimeAdapter", "MaterializedTabularInputView", "MapReduceAlgorithm", - "MetricPlan", - "OptimizationPlan", + "RayTorchAdapter", "ParallelEnsembleAlgorithm", "PortableRuntimeAdapter", "Predictable", "RuntimeExecutionEnvelope", + "TorchRuntimePreflight", "PreparedInput", "ResolvedInputLease", "RuntimeInputBinding", "TabularBatchInputView", "Transformable", - "TorchTrainingRecipe", - "TrainingRecipeV2", - "TrainingStepResult", + "TorchArtifactContext", + "TorchArtifactPlan", + "TorchBatch", + "TorchBatchContext", + "TorchBuildContext", + "TorchCheckpointContext", + "TorchMetricPlan", + "TorchModuleSet", + "TorchOptimizationPlan", + "TorchRecipe", + "TorchRuntimeContext", + "TorchStageContext", + "TorchStepContext", + "TorchStepResult", + "TorchWorkerCheckpointContext", "WorkerInputAdapter", "WorkerInputPayload", "WorkerInputPayloadSet", diff --git a/src/tributo/algorithms/spi/execution.py b/src/tributo/algorithms/spi/execution.py index 615181c..defbbdd 100644 --- a/src/tributo/algorithms/spi/execution.py +++ b/src/tributo/algorithms/spi/execution.py @@ -16,6 +16,7 @@ WorkerExecutionResult, ) from tributo.algorithms.api.distribution import StateField +from tributo.algorithms.api.torch_runtime import TorchPreflightLease from tributo.algorithms.spi.input import WorkerInputPayload, WorkerInputPayloadSet from tributo.util.annotations import PublicAPI @@ -67,6 +68,7 @@ def __post_init__(self) -> None: RuntimeTopology.DATA_PARALLEL, RuntimeTopology.RAY_MAP_REDUCE, RuntimeTopology.RAY_ITERATIVE_OPTIMIZATION, + RuntimeTopology.RAY_TRAIN_TORCH, } else 1 ) @@ -133,6 +135,22 @@ def runtime_id(self) -> str: ... def execute(self, envelope: RuntimeExecutionEnvelope) -> WorkerExecutionResult: ... +@PublicAPI(stability="alpha") +@runtime_checkable +class TorchRuntimePreflight(Protocol): + """Torch-only preflight surface kept out of the generic Runtime SPI.""" + + @property + def runtime_id(self) -> str: ... + + def preflight( + self, + plan: ResolvedAlgorithmPlan, + run_id: str, + invocation_id: str, + ) -> TorchPreflightLease: ... + + @PublicAPI(stability="alpha") class CollectiveAlgorithm(ABC): """Required surface for iterative Ray Train collective algorithms.""" @@ -421,5 +439,6 @@ def retry_safe(self) -> bool: "PortableRuntimeAdapter", "Predictable", "RuntimeExecutionEnvelope", + "TorchRuntimePreflight", "Transformable", ] diff --git a/src/tributo/algorithms/spi/torch.py b/src/tributo/algorithms/spi/torch.py index 9b5d5df..76653b5 100644 --- a/src/tributo/algorithms/spi/torch.py +++ b/src/tributo/algorithms/spi/torch.py @@ -1,126 +1,357 @@ -"""Narrow PyTorch recipe contract lowered to the Ray Train collective runtime.""" +"""Versioned PyTorch Recipe and Ray Torch Adapter contracts.""" from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Callable, Mapping +from collections.abc import Mapping from dataclasses import dataclass, field from types import MappingProxyType -from typing import Any - +from typing import Any, cast + +from tributo._common.immutable import deep_freeze +from tributo.algorithms.api.torch_runtime import ( + TorchCheckpointRef, + TorchCompositeLossContribution, + TorchLossContribution, + TorchMetricContribution, + TorchStageRunIdentity, + TorchStepLoss, +) from tributo.util.annotations import PublicAPI @PublicAPI(stability="alpha") -class TorchTrainingRecipe(ABC): - """Define model semantics without owning the distributed worker loop. +@dataclass(frozen=True) +class TorchRuntimeContext: + """Invocation-level context supplied by the Core Torch Runtime.""" + + algorithm_config: Mapping[str, Any] + implementation_id: str + world_size: int + policy_digest: str + execution_plan_digest: str + run_identity: TorchStageRunIdentity | None = None + + input_bindings: Mapping[str, object] = field(default_factory=dict) + output_config: Mapping[str, object] = field(default_factory=dict) + recovery_identity: Mapping[str, object] | None = None + input_binding_digest: str | None = None + state_layout: str = "replicated" + adapter_identity: str | None = None + resume_supported: bool = True + same_world_size_resume: bool | None = True + torch_runtime_api_version: int = 1 - A default recipe author implements only the four factory methods. Tributo - combines their products with Ray Train, Ray Data, and PyTorch DDP. The - optional ``forward`` and ``compute_loss`` methods are the bounded advanced - escape hatches; data sharding, reporting, checkpointing, and Bundle - publication remain framework-owned infrastructure. Multi-worker exact - coverage can invoke ``forward`` with a zero-row batch after another rank - exhausts its shard, so recipe models must preserve their output contract - for an empty leading batch dimension. - """ + def __post_init__(self) -> None: + if not isinstance(self.algorithm_config, Mapping): + raise ValueError("TorchRuntimeContext algorithm_config must be a mapping") + if self.torch_runtime_api_version != 1: + raise ValueError("TorchRuntimeContext torch_runtime_api_version must be 1") + if not isinstance(self.implementation_id, str) or not self.implementation_id: + raise ValueError("TorchRuntimeContext implementation_id is required") + if ( + not isinstance(self.world_size, int) + or isinstance(self.world_size, bool) + or self.world_size < 1 + ): + raise ValueError("TorchRuntimeContext world_size must be positive") + if not isinstance(self.policy_digest, str) or len(self.policy_digest) != 64: + raise ValueError("TorchRuntimeContext policy_digest is required") + if ( + not isinstance(self.execution_plan_digest, str) + or len(self.execution_plan_digest) != 64 + ): + raise ValueError("TorchRuntimeContext execution_plan_digest is required") + if self.run_identity is not None and not isinstance( + self.run_identity, TorchStageRunIdentity + ): + raise ValueError("TorchRuntimeContext run_identity is invalid") + if self.input_binding_digest is not None and ( + len(self.input_binding_digest) != 64 + or any(char not in "0123456789abcdef" for char in self.input_binding_digest) + ): + raise ValueError("TorchRuntimeContext input_binding_digest is invalid") + if self.state_layout not in {"replicated", "component", "sharded"}: + raise ValueError("TorchRuntimeContext state_layout is invalid") + if not isinstance(self.resume_supported, bool): + raise ValueError("TorchRuntimeContext resume_supported must be boolean") + if self.same_world_size_resume is not None and not isinstance( + self.same_world_size_resume, bool + ): + raise ValueError( + "TorchRuntimeContext same_world_size_resume must be boolean" + ) + object.__setattr__(self, "algorithm_config", deep_freeze(self.algorithm_config)) + object.__setattr__(self, "input_bindings", deep_freeze(self.input_bindings)) + object.__setattr__(self, "output_config", deep_freeze(self.output_config)) + if self.recovery_identity is not None: + object.__setattr__( + self, + "recovery_identity", + deep_freeze(self.recovery_identity), + ) + + def to_dict(self) -> dict[str, object]: + payload: dict[str, object] = { + "algorithm_config": dict(self.algorithm_config), + "implementation_id": self.implementation_id, + "world_size": self.world_size, + "policy_digest": self.policy_digest, + "execution_plan_digest": self.execution_plan_digest, + "run_identity": self.run_identity.to_dict() if self.run_identity else None, + "input_bindings": dict(self.input_bindings), + "output_config": dict(self.output_config), + "recovery_identity": dict(self.recovery_identity) + if self.recovery_identity + else None, + "input_binding_digest": self.input_binding_digest, + "state_layout": self.state_layout, + "adapter_identity": self.adapter_identity, + "resume_supported": self.resume_supported, + "torch_runtime_api_version": self.torch_runtime_api_version, + } + if self.same_world_size_resume is not None: + payload["same_world_size_resume"] = self.same_world_size_resume + return payload - api_version = 1 - @abstractmethod - def model_factory(self, config: Mapping[str, Any]) -> object: - """Build one worker-local ``torch.nn.Module`` from model config.""" +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchStageContext: + """Stage-aware context derived exclusively from ``TorchPolicy``.""" + + runtime: TorchRuntimeContext + stage_id: str + stage_index: int + is_final: bool + input_roles: tuple[str, ...] + predecessor_stage_id: str | None = None + predecessor_checkpoint_descriptor: Mapping[str, Any] | None = None + metric_mapping: Mapping[str, str] = field(default_factory=dict) + checkpoint_required: bool = True + checkpoint_interval_windows: int = 1 - @abstractmethod - def loss_factory(self, config: Mapping[str, Any]) -> object: - """Build a scalar batch-mean loss callable from loss config.""" + def __post_init__(self) -> None: + if not isinstance(self.runtime, TorchRuntimeContext): + raise ValueError("TorchStageContext runtime is required") + if not isinstance(self.stage_id, str) or not self.stage_id: + raise ValueError("TorchStageContext stage_id is required") + if ( + not isinstance(self.stage_index, int) + or isinstance(self.stage_index, bool) + or self.stage_index < 0 + ): + raise ValueError("TorchStageContext stage_index must be non-negative") + if not isinstance(self.is_final, bool): + raise ValueError("TorchStageContext is_final must be boolean") + object.__setattr__(self, "input_roles", tuple(self.input_roles)) + if any( + not isinstance(name, str) + or not name + or not isinstance(target, str) + or not target + for name, target in self.metric_mapping.items() + ): + raise ValueError("TorchStageContext metric_mapping is malformed") + if len(set(self.metric_mapping.values())) != len(self.metric_mapping): + raise ValueError("TorchStageContext metric_mapping targets must be unique") + if not isinstance(self.checkpoint_required, bool): + raise ValueError("TorchStageContext checkpoint_required must be boolean") + if ( + not isinstance(self.checkpoint_interval_windows, int) + or isinstance(self.checkpoint_interval_windows, bool) + or self.checkpoint_interval_windows < 1 + ): + raise ValueError( + "TorchStageContext checkpoint_interval_windows must be positive" + ) + object.__setattr__(self, "metric_mapping", deep_freeze(self.metric_mapping)) + if self.predecessor_checkpoint_descriptor is not None: + if any( + key in {"locator", "checkpoint_locator", "path", "credential"} + for key in self.predecessor_checkpoint_descriptor + ): + raise ValueError( + "TorchStageContext cannot expose checkpoint locator or credentials" + ) + object.__setattr__( + self, + "predecessor_checkpoint_descriptor", + deep_freeze(self.predecessor_checkpoint_descriptor), + ) + + def to_dict(self) -> dict[str, object]: + return { + "runtime": self.runtime.to_dict(), + "stage_id": self.stage_id, + "stage_index": self.stage_index, + "is_final": self.is_final, + "input_roles": list(self.input_roles), + "predecessor_stage_id": self.predecessor_stage_id, + "predecessor_checkpoint_descriptor": dict( + self.predecessor_checkpoint_descriptor + ) + if self.predecessor_checkpoint_descriptor + else None, + "metric_mapping": dict(self.metric_mapping), + "checkpoint_required": self.checkpoint_required, + "checkpoint_interval_windows": self.checkpoint_interval_windows, + } + + @classmethod + def from_dict(cls, value: Mapping[str, object]) -> "TorchStageContext": + runtime_value = value.get("runtime") + if not isinstance(runtime_value, Mapping): + raise ValueError("TorchStageContext runtime payload is invalid") + identity_value = runtime_value.get("run_identity") + resume_supported = runtime_value.get("resume_supported", True) + runtime = TorchRuntimeContext( + algorithm_config=runtime_value.get("algorithm_config", {}), + implementation_id=runtime_value["implementation_id"], + world_size=runtime_value["world_size"], + policy_digest=runtime_value["policy_digest"], + execution_plan_digest=runtime_value["execution_plan_digest"], + run_identity=( + TorchStageRunIdentity.from_dict(identity_value) + if isinstance(identity_value, Mapping) + else None + ), + input_bindings=runtime_value.get("input_bindings", {}), + output_config=runtime_value.get("output_config", {}), + recovery_identity=runtime_value.get("recovery_identity"), + input_binding_digest=runtime_value.get("input_binding_digest"), + state_layout=runtime_value.get("state_layout", "replicated"), + adapter_identity=runtime_value.get("adapter_identity"), + resume_supported=resume_supported, + same_world_size_resume=runtime_value.get( + "same_world_size_resume", True if resume_supported else None + ), + torch_runtime_api_version=runtime_value.get("torch_runtime_api_version", 1), + ) + return cls( + runtime=runtime, + stage_id=cast(str, value["stage_id"]), + stage_index=cast(int, value["stage_index"]), + is_final=cast(bool, value["is_final"]), + input_roles=tuple(cast(tuple[str, ...], value.get("input_roles", ()))), + predecessor_stage_id=cast(str | None, value.get("predecessor_stage_id")), + predecessor_checkpoint_descriptor=cast( + Mapping[str, Any] | None, + value.get("predecessor_checkpoint_descriptor"), + ), + metric_mapping=cast(Mapping[str, str], value.get("metric_mapping", {})), + checkpoint_required=cast(bool, value.get("checkpoint_required", True)), + checkpoint_interval_windows=cast( + int, value.get("checkpoint_interval_windows", 1) + ), + ) - @abstractmethod - def optimizer_factory( - self, - model: object, - config: Mapping[str, Any], - ) -> object: - """Build an optimizer for the unwrapped model from optimizer config.""" - @abstractmethod - def metric_factories( - self, - config: Mapping[str, Any], - ) -> Mapping[str, Callable[[object, object], object]]: - """Build metric callables keyed by declared metric name.""" +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchCheckpointContext: + """Control metadata passed beside an actual Ray Train result/checkpoint.""" - def forward(self, model: object, features: object) -> object: - """Invoke the default one-tensor model signature.""" - if not callable(model): - raise TypeError("recipe model must be callable") - return model(features) + stage: TorchStageContext + run_id: str + invocation_id: str + checkpoint_owner: str - def compute_loss( - self, - loss: object, - predictions: object, - targets: object, - ) -> object: - """Invoke the default ``loss(predictions, targets)`` signature.""" - if not callable(loss): - raise TypeError("recipe loss must be callable") - return loss(predictions, targets) + def __post_init__(self) -> None: + if not isinstance(self.stage, TorchStageContext): + raise ValueError("TorchCheckpointContext stage is required") + for name in ("run_id", "invocation_id", "checkpoint_owner"): + if not isinstance(getattr(self, name), str) or not getattr(self, name): + raise ValueError(f"TorchCheckpointContext {name} is required") @PublicAPI(stability="alpha") @dataclass(frozen=True) -class OptimizationPlan: - """Optimizer and bounded loop controls selected by one RecipeV2.""" +class TorchWorkerCheckpointContext: + """Worker-local, non-serializable view of a selected checkpoint.""" - optimizer: object - scheduler: object | None = None - gradient_accumulation_steps: int = 1 - max_gradient_norm: float | None = None + stage: TorchStageContext + source: str + checkpoint: TorchCheckpointRef | None = None def __post_init__(self) -> None: - if ( - not isinstance(self.gradient_accumulation_steps, int) - or isinstance(self.gradient_accumulation_steps, bool) - or self.gradient_accumulation_steps < 1 + if not isinstance(self.stage, TorchStageContext): + raise ValueError("TorchWorkerCheckpointContext stage is required") + if self.source not in { + "none", + "ray_failure_retry", + "stage_dependency", + "cross_run_initial_recovery", + }: + raise ValueError("TorchWorkerCheckpointContext source is invalid") + if self.source == "none" and self.checkpoint is not None: + raise ValueError("empty Torch checkpoint source cannot carry a checkpoint") + if self.checkpoint is not None and not isinstance( + self.checkpoint, TorchCheckpointRef ): - raise ValueError("gradient_accumulation_steps must be positive") - if self.max_gradient_norm is not None and ( - not isinstance(self.max_gradient_norm, (int, float)) - or isinstance(self.max_gradient_norm, bool) - or float(self.max_gradient_norm) <= 0 - ): - raise ValueError("max_gradient_norm must be positive when provided") + raise ValueError("TorchWorkerCheckpointContext checkpoint is invalid") @PublicAPI(stability="alpha") @dataclass(frozen=True) -class MetricPlan: - """Bounded metric callables keyed by descriptor-declared identities.""" - - factories: Mapping[str, Callable[[object, object], object]] = field( - default_factory=dict - ) +class TorchBatch: + """Typed algorithm batch preserving named inputs and exact row coverage.""" + + positional: tuple[object, ...] = () + keyword: Mapping[str, object] = field(default_factory=dict) + targets: object | None = None + weights: object | None = None + local_rows: int = 0 + coverage_counts: Mapping[str, int] = field(default_factory=dict) def __post_init__(self) -> None: - if any( - not isinstance(name, str) or not name or not callable(metric) - for name, metric in self.factories.items() - ): - raise ValueError("MetricPlan requires named callable metrics") + if self.local_rows < 0 or isinstance(self.local_rows, bool): + raise ValueError("TorchBatch local_rows must be non-negative") + for name, count in self.coverage_counts.items(): + if ( + not isinstance(name, str) + or not name + or not isinstance(count, int) + or isinstance(count, bool) + or count < 0 + ): + raise ValueError( + "TorchBatch coverage counts must be non-negative integers" + ) + object.__setattr__(self, "keyword", MappingProxyType(dict(self.keyword))) + object.__setattr__( + self, + "coverage_counts", + MappingProxyType(dict(self.coverage_counts)), + ) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchStepContext: + """Per-step context with no access to Core lifecycle internals.""" + + stage: TorchStageContext + window_index: int + micro_batch_index: int @PublicAPI(stability="alpha") @dataclass(frozen=True) -class TrainingStepResult: - """One algorithm-owned forward/loss result consumed by the Core loop.""" +class TorchStepResult: + """Forward outputs and an explicit ordinary/composite loss contribution.""" - predictions: object - loss: object + outputs: Mapping[str, object] + loss: TorchStepLoss coverage_counts: Mapping[str, int] = field(default_factory=dict) + metrics: Mapping[str, TorchMetricContribution] = field(default_factory=dict) def __post_init__(self) -> None: - normalized: dict[str, int] = {} + if not isinstance( + self.loss, (TorchLossContribution, TorchCompositeLossContribution) + ): + raise ValueError("TorchStepResult requires a typed TorchStepLoss") + object.__setattr__(self, "outputs", MappingProxyType(dict(self.outputs))) for name, count in self.coverage_counts.items(): if ( not isinstance(name, str) @@ -130,78 +361,292 @@ def __post_init__(self) -> None: or count < 0 ): raise ValueError( - "TrainingStepResult coverage counts require named non-negative " - "integers" + "TorchStepResult coverage counts must be non-negative integers" ) - normalized[name] = count - object.__setattr__(self, "coverage_counts", MappingProxyType(normalized)) + if any( + not isinstance(value, TorchMetricContribution) + for value in self.metrics.values() + ): + raise ValueError( + "TorchStepResult metrics must be TorchMetricContribution values" + ) + object.__setattr__( + self, "coverage_counts", MappingProxyType(dict(self.coverage_counts)) + ) + object.__setattr__(self, "metrics", MappingProxyType(dict(self.metrics))) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchModuleSet: + """Named worker modules produced by a :class:`TorchRecipe`.""" + + modules: Mapping[str, object] + + def __post_init__(self) -> None: + values = dict(self.modules) + if "model" not in values or "loss" not in values: + raise ValueError("TorchModuleSet requires model and loss modules") + if any(not isinstance(name, str) or not name for name in values): + raise ValueError("TorchModuleSet module names must be non-empty") + object.__setattr__(self, "modules", MappingProxyType(values)) + + def __getitem__(self, name: str) -> object: + return self.modules[name] + + def get(self, name: str, default: object | None = None) -> object | None: + return self.modules.get(name, default) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchArtifactPlan: + """Typed export declaration consumed by Core BundleExportService.""" + + source_kind: str + input_signature: tuple[Mapping[str, object], ...] + output_signature: tuple[Mapping[str, object], ...] + targets: tuple[Mapping[str, object], ...] + roles: Mapping[str, str] + required: bool = True + + def __post_init__(self) -> None: + if not isinstance(self.source_kind, str) or not self.source_kind: + raise ValueError("TorchArtifactPlan source_kind is required") + for name, values in ( + ("input_signature", self.input_signature), + ("output_signature", self.output_signature), + ("targets", self.targets), + ): + if any(not isinstance(value, Mapping) for value in values): + raise ValueError(f"TorchArtifactPlan {name} entries must be mappings") + normalized = tuple(dict(value) for value in values) + object.__setattr__(self, name, normalized) + target_names: list[str] = [] + for target in self.targets: + target_name = target.get("name") + if not isinstance(target_name, str) or not target_name: + raise ValueError("TorchArtifactPlan targets require named targets") + target_names.append(target_name) + if len(set(target_names)) != len(target_names): + raise ValueError("TorchArtifactPlan target names must be unique") + for signature_name in ("input_signature", "output_signature"): + for field_value in getattr(self, signature_name): + if ( + not isinstance(field_value.get("name"), str) + or not field_value["name"] + or not isinstance(field_value.get("dtype"), str) + or not field_value["dtype"] + or not isinstance(field_value.get("shape", ()), (list, tuple)) + ): + raise ValueError(f"TorchArtifactPlan {signature_name} is malformed") + if not self.output_signature: + raise ValueError("TorchArtifactPlan requires an output signature") + if not isinstance(self.required, bool): + raise ValueError("TorchArtifactPlan required must be boolean") + roles = dict(self.roles) + if any( + not isinstance(name, str) + or not name + or not isinstance(target, str) + or not target + for name, target in roles.items() + ): + raise ValueError("TorchArtifactPlan roles must map names to targets") + if any(target not in target_names for target in roles.values()): + raise ValueError("TorchArtifactPlan roles must reference declared targets") + object.__setattr__(self, "roles", MappingProxyType(roles)) + + def to_dict(self) -> dict[str, object]: + return { + "source_kind": self.source_kind, + "input_signature": [dict(value) for value in self.input_signature], + "output_signature": [dict(value) for value in self.output_signature], + "targets": [dict(value) for value in self.targets], + "roles": dict(self.roles), + "required": self.required, + } + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchBuildContext: + """Model-construction context for a Core-managed Recipe.""" + + runtime: TorchRuntimeContext + stage: TorchStageContext + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchBatchContext: + """Batch adaptation context with named input role information.""" + + stage: TorchStageContext + input_roles: tuple[str, ...] = () + feature_names: tuple[str, ...] = () + label_name: str | None = None + weight_name: str | None = None + + def __post_init__(self) -> None: + roles = tuple(self.input_roles or self.stage.input_roles) + if not roles or any(not isinstance(role, str) or not role for role in roles): + raise ValueError("TorchBatchContext input_roles must be non-empty") + if len(set(roles)) != len(roles): + raise ValueError("TorchBatchContext input_roles must be unique") + object.__setattr__(self, "input_roles", roles) @PublicAPI(stability="alpha") -class TrainingRecipeV2(ABC): - """Define PyTorch mathematics while Core owns the distributed loop.""" +@dataclass(frozen=True) +class TorchOptimizationPlan: + """Optimizer and accumulation controls selected by one Recipe.""" - api_version = 2 + optimizer: object + scheduler: object | None = None + gradient_accumulation_steps: int = 1 + max_gradient_norm: float | None = None + + def __post_init__(self) -> None: + if ( + not isinstance(self.gradient_accumulation_steps, int) + or isinstance(self.gradient_accumulation_steps, bool) + or self.gradient_accumulation_steps < 1 + ): + raise ValueError("gradient_accumulation_steps must be positive") + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchMetricPlan: + """Metric names and reducer identities for one Torch Policy.""" + + reducers: Mapping[str, str] + + def __post_init__(self) -> None: + normalized = dict(self.reducers) + if any(not isinstance(k, str) or not k for k in normalized): + raise ValueError("TorchMetricPlan reducer names must be non-empty") + object.__setattr__(self, "reducers", MappingProxyType(normalized)) + + +@PublicAPI(stability="alpha") +@dataclass(frozen=True) +class TorchArtifactContext: + """Context used to construct a typed artifact plan.""" + + stage: TorchStageContext + checkpoint: TorchCheckpointRef | None = None + + def __post_init__(self) -> None: + if not isinstance(self.stage, TorchStageContext): + raise ValueError("TorchArtifactContext stage is required") + + +@PublicAPI(stability="alpha") +class TorchRecipe(ABC): + """Core-owned PyTorch training hooks with no Runtime lifecycle control.""" + + api_version = 1 @abstractmethod - def build_modules(self, config: Mapping[str, Any]) -> Mapping[str, object]: - """Build at least ``model`` and ``loss`` modules.""" + def build_modules(self, context: TorchBuildContext) -> TorchModuleSet: ... @abstractmethod - def batch_adapter( - self, - batch: object, - *, - feature_names: tuple[str, ...], - label_name: str | None, - weight_name: str | None, - config: Mapping[str, Any], - ) -> tuple[object, object, object | None, int]: - """Convert a Ray batch into features, targets, weights, and row count.""" + def adapt_batch(self, batch: object, context: TorchBatchContext) -> TorchBatch: ... @abstractmethod def training_step( self, - modules: Mapping[str, object], - features: object, - targets: object, - weights: object | None, - config: Mapping[str, Any], - ) -> TrainingStepResult: - """Compute predictions and one scalar batch-mean training loss.""" + modules: TorchModuleSet, + batch: TorchBatch, + context: TorchStepContext, + ) -> TorchStepResult: ... @abstractmethod def validation_step( self, - modules: Mapping[str, object], - features: object, - targets: object, - weights: object | None, - config: Mapping[str, Any], - ) -> TrainingStepResult: - """Compute predictions and one scalar validation loss.""" + modules: TorchModuleSet, + batch: TorchBatch, + context: TorchStepContext, + ) -> TorchStepResult: ... + + @abstractmethod + def configure_optimizers( + self, + modules: TorchModuleSet, + context: TorchBuildContext, + ) -> TorchOptimizationPlan: ... + + @abstractmethod + def metric_plan(self, context: TorchRuntimeContext) -> TorchMetricPlan: ... + + @abstractmethod + def artifact_plan(self, context: TorchArtifactContext) -> TorchArtifactPlan: ... + + +@PublicAPI(stability="alpha") +class RayTorchAdapter(ABC): + """Framework-native Torch hooks driven by the Core Ray Train Runtime.""" + + api_version = 1 + + @abstractmethod + def validate_environment(self, context: TorchRuntimeContext) -> None: ... + + @abstractmethod + def bind_datasets( + self, + datasets: Mapping[str, object], + context: TorchStageContext, + ) -> Mapping[str, object]: ... + + @abstractmethod + def worker_config(self, context: TorchStageContext) -> Mapping[str, object]: ... @abstractmethod - def optimization_plan( + def train_loop_per_worker( self, - model: object, - config: Mapping[str, Any], - ) -> OptimizationPlan: - """Build optimizer, optional scheduler, accumulation, and clipping.""" + worker_config: Mapping[str, object], + checkpoint_context: TorchWorkerCheckpointContext, + ) -> None: ... @abstractmethod - def metric_plan(self, config: Mapping[str, Any]) -> MetricPlan: - """Build bounded metrics consumed by Core collective reduction.""" + def checkpoint_source( + self, + result: object, + context: TorchCheckpointContext, + ) -> object: ... @abstractmethod - def checkpoint_codec(self) -> object: - """Return a codec for algorithm-specific checkpoint payload state.""" + def metric_plan(self, context: TorchRuntimeContext) -> TorchMetricPlan: ... + + @abstractmethod + def artifact_plan(self, context: TorchArtifactContext) -> TorchArtifactPlan: ... + + @abstractmethod + def open_export_source( + self, + checkpoint_ref: TorchCheckpointRef, + artifact_context: TorchArtifactContext, + ) -> Any: ... __all__ = [ - "MetricPlan", - "OptimizationPlan", - "TorchTrainingRecipe", - "TrainingRecipeV2", - "TrainingStepResult", + "RayTorchAdapter", + "TorchArtifactContext", + "TorchArtifactPlan", + "TorchBatch", + "TorchBatchContext", + "TorchBuildContext", + "TorchCheckpointContext", + "TorchMetricPlan", + "TorchModuleSet", + "TorchOptimizationPlan", + "TorchRecipe", + "TorchRuntimeContext", + "TorchStageContext", + "TorchStepContext", + "TorchStepResult", + "TorchWorkerCheckpointContext", ] diff --git a/src/tributo/inference/kernel.py b/src/tributo/inference/kernel.py index 3c7c531..edd4f6e 100644 --- a/src/tributo/inference/kernel.py +++ b/src/tributo/inference/kernel.py @@ -153,6 +153,14 @@ def _build_input_tensor( "scalar single-column input must be a one-dimensional batch column" ) tensor = arrays[0] + elif len(arrays) == 1 and arrays[0].dtype == object: + # Arrow/Parquet may decode a vector-valued column as either a one- or + # multi-dimensional object array whose rows contain nested arrays. + # Stack rows before dtype conversion so the declared tensor rank is + # preserved instead of treating the object dimension as a feature axis. + tensor = _stack_nested_object_array(arrays[0]) + if tensor.dtype == object: + tensor = np.column_stack(arrays) elif len(arrays) == 1 and arrays[0].ndim > 1: tensor = arrays[0] else: @@ -173,6 +181,20 @@ def _build_input_tensor( return np.asarray(tensor) +def _stack_nested_object_array(value: object) -> np.ndarray: + """Materialize homogeneous nested object arrays while preserving row rank.""" + array = np.asarray(value) + if array.dtype != object: + return array + try: + nested = [_stack_nested_object_array(item) for item in array.tolist()] + if nested: + return np.stack(nested) + except (TypeError, ValueError): + pass + return array + + def _tensor_row_count(tensors: dict[str, np.ndarray], *, kind: str) -> int: row_count: int | None = None for name, tensor in tensors.items(): diff --git a/src/tributo/integrations/algorithm_inputs/ingestion.py b/src/tributo/integrations/algorithm_inputs/ingestion.py index 78c25b3..6e010a2 100644 --- a/src/tributo/integrations/algorithm_inputs/ingestion.py +++ b/src/tributo/integrations/algorithm_inputs/ingestion.py @@ -436,6 +436,24 @@ def bind( "ingestion runtime adapter requires a typed Gateway handle" ) binding = lease.binding or plan.primary_input_binding + if plan.runtime.topology is RuntimeTopology.RAY_TRAIN_TORCH: + if not isinstance(lease.handle, RayDataHandle): + raise AlgorithmInputError( + "Ray Train Torch input requires RayDataHandle; no implicit " + "Daft-to-Ray conversion is permitted" + ) + return RuntimeInputBinding( + tuple( + WorkerInputPayload( + input_name=binding.name, + binding=binding, + value=lease.handle, + partition_index=rank, + partition_count=plan.runtime.worker_count, + ) + for rank in range(plan.runtime.worker_count) + ) + ) if plan.runtime.topology in { RuntimeTopology.DATA_PARALLEL, RuntimeTopology.RAY_MAP_REDUCE, diff --git a/src/tributo/integrations/algorithm_runtimes/collective.py b/src/tributo/integrations/algorithm_runtimes/collective.py index cc054b2..d2259c6 100644 --- a/src/tributo/integrations/algorithm_runtimes/collective.py +++ b/src/tributo/integrations/algorithm_runtimes/collective.py @@ -38,9 +38,6 @@ RAY_TRAIN_COLLECTIVE_RUNTIME_ID = FORMAL_DISTRIBUTED_STRATEGY_CONTRACTS[ DistributionStrategy.RAY_TRAIN_COLLECTIVE ].runtime_id -RAY_TRAIN_RECIPE_V2_RUNTIME_ID = FORMAL_DISTRIBUTED_STRATEGY_CONTRACTS[ - DistributionStrategy.RAY_TRAIN_RECIPE_V2 -].runtime_id def _ray_run_name(algorithm: str, run_id: str) -> str: @@ -90,12 +87,9 @@ def _collective_execution_result( def _load_algorithm(envelope: RuntimeExecutionEnvelope) -> CollectiveAlgorithm: plan = envelope.plan spec = plan.distribution_spec - if spec is None or spec.strategy not in { - DistributionStrategy.RAY_TRAIN_COLLECTIVE, - DistributionStrategy.RAY_TRAIN_RECIPE_V2, - }: + if spec is None or spec.strategy is not DistributionStrategy.RAY_TRAIN_COLLECTIVE: raise AlgorithmConfigurationError( - "collective runtime requires a collective or RecipeV2 DistributionSpec" + "collective runtime requires a RAY_TRAIN_COLLECTIVE DistributionSpec" ) _validate_module_digest( plan.implementation.implementation_ref, @@ -424,14 +418,4 @@ def execute(self, envelope: RuntimeExecutionEnvelope) -> WorkerExecutionResult: prepared.close() -@DeveloperAPI -class RayTrainRecipeV2Runtime(RayTrainCollectiveRuntime): - """Execute TrainingRecipeV2 through the same Core-owned Ray Train loop.""" - - @property - def runtime_id(self) -> str: - """Return the dedicated RecipeV2 runtime identity.""" - return RAY_TRAIN_RECIPE_V2_RUNTIME_ID - - -__all__ = ["RayTrainCollectiveRuntime", "RayTrainRecipeV2Runtime"] +__all__ = ["RayTrainCollectiveRuntime"] diff --git a/src/tributo/integrations/algorithm_runtimes/portable_metrics.py b/src/tributo/integrations/algorithm_runtimes/portable_metrics.py index e2d2891..38fdd08 100644 --- a/src/tributo/integrations/algorithm_runtimes/portable_metrics.py +++ b/src/tributo/integrations/algorithm_runtimes/portable_metrics.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Collection, Mapping from typing import Any from tributo.algorithms.api import ( @@ -26,7 +26,11 @@ ) -def portable_fit_only_metrics(metrics: Mapping[Any, Any]) -> dict[str, Any]: +def portable_fit_only_metrics( + metrics: Mapping[Any, Any], + *, + extra_evidence_names: Collection[str] = (), +) -> dict[str, Any]: """Return user metrics after one fail-closed portable-value validation. Runtime evidence fields are intentionally removed before validation. Any @@ -37,10 +41,11 @@ def portable_fit_only_metrics(metrics: Mapping[Any, Any]) -> dict[str, Any]: raise AlgorithmExecutionError( "FIT_ONLY user metrics must be provided as a mapping" ) + evidence_names = FIT_ONLY_EVIDENCE_METRIC_NAMES | frozenset(extra_evidence_names) user_metrics = { name: value for name, value in metrics.items() - if not (isinstance(name, str) and name in FIT_ONLY_EVIDENCE_METRIC_NAMES) + if not (isinstance(name, str) and name in evidence_names) } if any(not isinstance(name, str) for name in user_metrics): raise AlgorithmExecutionError("FIT_ONLY user metric names must be strings") diff --git a/src/tributo/integrations/algorithm_runtimes/ray_data_config.py b/src/tributo/integrations/algorithm_runtimes/ray_data_config.py index e367bfc..695e6e7 100644 --- a/src/tributo/integrations/algorithm_runtimes/ray_data_config.py +++ b/src/tributo/integrations/algorithm_runtimes/ray_data_config.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Mapping from typing import TYPE_CHECKING, Any from ray.train import DataConfig @@ -72,4 +73,17 @@ def configure( return output -__all__ = ["ExactCoverageDataConfig"] +@DeveloperAPI +class TorchRoleDataConfig(ExactCoverageDataConfig): + """Route Torch roles independently: split_exact roles shard, replicate roles do not.""" + + def __init__(self, routes: Mapping[str, object], **kwargs: Any) -> None: + split_roles = [ + role + for role, route in routes.items() + if getattr(route, "mode", None) == "split_exact" + ] + super().__init__(datasets_to_split=split_roles, **kwargs) + + +__all__ = ["ExactCoverageDataConfig", "TorchRoleDataConfig"] diff --git a/src/tributo/integrations/algorithm_runtimes/ray_train_torch.py b/src/tributo/integrations/algorithm_runtimes/ray_train_torch.py new file mode 100644 index 0000000..b452a5b --- /dev/null +++ b/src/tributo/integrations/algorithm_runtimes/ray_train_torch.py @@ -0,0 +1,4028 @@ +"""Core-owned Ray Train Torch runtime. + +The module is intentionally the only Runtime implementation selected by +``DistributionStrategy.RAY_TRAIN_TORCH``. It owns Trainer construction and +Stage ordering; algorithm Wheels only provide a ``TorchRecipe`` or +``RayTorchAdapter`` implementation. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import math +import os +import tempfile +from collections.abc import Mapping +from contextlib import contextmanager, nullcontext +from pathlib import Path +from typing import Any, Generator, cast + +from tributo.algorithms.api import ( + AlgorithmConfigurationError, + AlgorithmExecutionError, + AlgorithmExecutionResult, + ComponentStageEvidence, + DistributionStrategy, + QualifiedReference, + ReplicatedTorchStateEvidence, + ResultPolicy, + TorchAccumulationWindow, + TorchBackwardContext, + TorchCheckpointDescriptor, + TorchCheckpointLocator, + TorchCheckpointProgress, + TorchCheckpointRef, + TorchCompositeGlobalState, + TorchCompositeLossContribution, + TorchExecutionEvidence, + TorchGlobalLossContext, + TorchGlobalLossReducer, + TorchGlobalLossReduction, + TorchLossContribution, + TorchMetricContribution, + TorchPreflightLease, + TorchPreflightTokenData, + TorchRankProgressStatistics, + TorchRecoveryEnvelope, + TorchRoleExecutionEvidence, + TorchRuntimeExecutionEnvelope, + TorchStageRunIdentity, + TorchWorkerControlEnvelope, + WorkerExecutionEvidence, + WorkerExecutionResult, + apply_torch_loss_backward, + claim_torch_run_directory, + describe_torch_checkpoint, + invoke_torch_global_loss_reducer, + report_torch_checkpoint, + torch_run_config_name, + validate_torch_retry_identity, +) +from tributo.algorithms.api.models import FORMAL_DISTRIBUTED_STRATEGY_CONTRACTS +from tributo.algorithms.core.worker import ( + _actual_environment_versions, + _load_reference, + _validate_module_digest, +) +from tributo.algorithms.spi import ( + PreparedInput, + RayTorchAdapter, + RuntimeExecutionEnvelope, + TorchBatch, + TorchBatchContext, + TorchBuildContext, + TorchCheckpointContext, + TorchMetricPlan, + TorchModuleSet, + TorchOptimizationPlan, + TorchRecipe, + TorchRuntimeContext, + TorchStageContext, + TorchStepContext, + TorchStepResult, + TorchWorkerCheckpointContext, +) +from tributo.integrations.algorithm_runtimes.portable_metrics import ( + portable_fit_only_metrics, +) +from tributo.util.annotations import DeveloperAPI + +RAY_TRAIN_TORCH_RUNTIME_ID = FORMAL_DISTRIBUTED_STRATEGY_CONTRACTS[ + DistributionStrategy.RAY_TRAIN_TORCH +].runtime_id +_CORE_RECIPE_LOOP_REF = ( + "tributo.integrations.algorithm_runtimes.ray_train_torch:" + "torch_recipe_train_loop_per_worker" +) +_CORE_ADAPTER_LOOP_REF = ( + "tributo.integrations.algorithm_runtimes.ray_train_torch:" + "ray_torch_adapter_train_loop_per_worker" +) +logger = logging.getLogger(__name__) + + +def _load_torch_implementation(plan: Any) -> TorchRecipe | RayTorchAdapter: + """Load a trusted recipe/adapter exactly once after preflight.""" + _validate_module_digest( + plan.implementation.implementation_ref, + plan.implementation.code_digest, + ) + implementation = _load_reference(plan.implementation.implementation_ref) + if not isinstance(implementation, type): + raise AlgorithmConfigurationError( + "Torch implementation reference must resolve to a class" + ) + if not issubclass(implementation, (TorchRecipe, RayTorchAdapter)): + raise AlgorithmConfigurationError( + "Torch implementation must subclass TorchRecipe or RayTorchAdapter" + ) + if getattr(implementation, "api_version", None) != 1: + raise AlgorithmConfigurationError( + "Torch implementation api_version must be exactly 1" + ) + try: + instance = implementation() + except TypeError as exc: + raise AlgorithmConfigurationError( + "Torch implementation must have a no-argument constructor" + ) from exc + return cast(TorchRecipe | RayTorchAdapter, instance) + + +def _policy(plan: Any) -> Any: + distribution = plan.distribution_spec + if ( + distribution is None + or distribution.strategy is not DistributionStrategy.RAY_TRAIN_TORCH + ): + raise AlgorithmConfigurationError( + "Ray Train Torch runtime requires RAY_TRAIN_TORCH" + ) + policy = distribution.policy + if not hasattr(policy, "execution_plan"): + raise AlgorithmConfigurationError("Torch runtime lost TorchPolicy") + return policy + + +def _torch_algorithm_context_config(plan: Any) -> dict[str, object]: + """Return only algorithm-owned config for Torch implementation contexts.""" + config = plan.algorithm_config + if not isinstance(config, Mapping): + raise AlgorithmConfigurationError("Torch algorithm config must be a mapping") + return { + str(key): value + for key, value in config.items() + if str(key) not in {"ray", "output"} + } + + +def _torch_input_bindings(plan: Any) -> dict[str, object]: + """Expose credential-free InputBinding metadata to Torch implementations.""" + return { + binding.name: binding.descriptor_payload() + for binding in plan.input_bindings.bindings + } + + +def _torch_output_config(plan: Any) -> dict[str, object]: + """Expose output options separately from algorithm worker configuration.""" + config = plan.algorithm_config + output = config.get("output", {}) if isinstance(config, Mapping) else {} + if not isinstance(output, Mapping): + raise AlgorithmConfigurationError("Torch output config must be a mapping") + return {str(key): value for key, value in output.items()} + + +_ADAPTER_CONFIG_BLOCKED_KEYS = frozenset( + { + "ray", + "output", + "path", + "uri", + "locator", + "storage_path", + "resume", + "resume_from", + "checkpoint", + "checkpoint_uri", + "checkpoint_locator", + "bundle_uri", + "credential", + "credentials", + "secret", + "secrets", + } +) +_TORCH_INTERNAL_METRIC_NAMES = frozenset( + { + "torch_evidence", + "checkpoint_descriptor", + "reducer_id", + "reducer_api_version", + "reducer_schema_id", + "reducer_code_digest", + "reducer_branch", + "reducer_evidence", + } +) + + +def _validate_adapter_worker_config( + value: object, *, path: str = "adapter_config" +) -> None: + """Reject Core paths, recovery handles and credentials in Adapter config.""" + if isinstance(value, Mapping): + for raw_key, nested in value.items(): + if not isinstance(raw_key, str): + raise AlgorithmConfigurationError( + "Adapter worker config keys must be strings" + ) + key = raw_key.casefold() + if key in _ADAPTER_CONFIG_BLOCKED_KEYS or key.endswith( + ("_path", "_uri", "_locator") + ): + raise AlgorithmConfigurationError( + f"Adapter worker config contains a Core-owned path field: {path}.{raw_key}" + ) + _validate_adapter_worker_config(nested, path=f"{path}.{raw_key}") + elif isinstance(value, (list, tuple)): + for index, nested in enumerate(value): + _validate_adapter_worker_config(nested, path=f"{path}[{index}]") + + +def _accumulate_metric_totals( + target: dict[str, list[float]], + contributions: Mapping[str, TorchMetricContribution], + *, + prefix: str = "", +) -> None: + """Accumulate typed metric contributions without interpreting their meaning.""" + for name, contribution in contributions.items(): + if not isinstance(name, str) or not isinstance( + contribution, TorchMetricContribution + ): + raise AlgorithmConfigurationError( + "Torch reducer metric contribution is invalid" + ) + key = f"{prefix}{name}" + totals = target.setdefault(key, [0.0, 0.0]) + totals[0] += contribution.numerator + totals[1] += contribution.normalizer + + +def _reduce_metric_totals( + totals: Mapping[str, list[float]], + reducers: Mapping[str, str], + *, + device: object, + dist: Any, + world_size: int, +) -> dict[str, float]: + """Apply the reducer declared for each metric with aligned collectives.""" + import torch + + names = set(totals) + if dist.is_available() and dist.is_initialized(): + gathered: list[object] = [None] * world_size + dist.all_gather_object(gathered, sorted(names)) + if any(value != sorted(names) for value in gathered): + raise AlgorithmExecutionError("Torch metric keys differ across ranks") + values: dict[str, float] = {} + for name in sorted(names): + numerator, normalizer = totals[name] + reducer_value = reducers.get(name) + if reducer_value is None and "_" in name: + reducer_value = reducers.get(name.split("_", 1)[1]) + if reducer_value is None and name.endswith("_loss"): + reducer_value = reducers.get("train_loss") + if reducer_value is None: + raise AlgorithmConfigurationError( + f"Torch metric {name!r} has no declared reducer" + ) + reducer = str(reducer_value) + if reducer in {"sum_count", "weighted_mean"}: + state = torch.tensor( + [float(numerator), float(normalizer)], + dtype=torch.float64, + device=device, + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(state, op=dist.ReduceOp.SUM) + if state[1].item() <= 0: + raise AlgorithmExecutionError( + f"Torch metric {name!r} has a zero global normalizer" + ) + values[name] = float(state[0].item() / state[1].item()) + continue + if reducer not in {"min", "max"}: + raise AlgorithmConfigurationError( + f"unsupported Torch metric reducer {reducer!r}" + ) + present = normalizer > 0 + sentinel = float("inf") if reducer == "min" else float("-inf") + state = torch.tensor( + float(numerator / normalizer) if present else sentinel, + dtype=torch.float64, + device=device, + ) + present_state = torch.tensor( + 1 if present else 0, dtype=torch.int64, device=device + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(present_state, op=dist.ReduceOp.MAX) + dist.all_reduce( + state, + op=torch.distributed.ReduceOp.MIN + if reducer == "min" + else torch.distributed.ReduceOp.MAX, + ) + if int(present_state.item()) == 0: + raise AlgorithmExecutionError( + f"Torch metric {name!r} has no global contribution" + ) + values[name] = float(state.item()) + return values + + +def _reducer_metadata(policy: Any) -> dict[str, object]: + """Load the qualified reducer metadata already attested by preflight.""" + if policy.global_loss_reducer_ref is None: + return {} + reference = QualifiedReference.parse(policy.global_loss_reducer_ref) + _validate_module_digest(reference, policy.global_loss_reducer_code_digest) + reducer = _load_reference(reference) + if isinstance(reducer, type): + reducer = reducer() + values = { + "reducer_id": getattr(reducer, "reducer_id", None), + "reducer_api_version": getattr(reducer, "api_version", None), + "reducer_schema_id": getattr(reducer, "component_schema_id", None), + "reducer_code_digest": getattr(reducer, "code_digest", None), + } + if ( + not isinstance(values["reducer_id"], str) + or values["reducer_api_version"] != policy.global_loss_reducer_api_version + or values["reducer_schema_id"] != policy.composite_loss_schema_id + or values["reducer_code_digest"] != policy.global_loss_reducer_code_digest + ): + raise AlgorithmConfigurationError("Torch global loss reducer metadata drifted") + return values + + +def policy_result_policy(plan: Any) -> ResultPolicy: + distribution = plan.distribution_spec + if distribution is None: + raise AlgorithmConfigurationError("Torch plan has no DistributionSpec") + return cast(ResultPolicy, distribution.result_policy) + + +def _identity( + plan: Any, run_id: str, invocation_id: str, stage_id: str +) -> TorchStageRunIdentity: + code_digest = plan.implementation.code_digest + if not isinstance(code_digest, str) or len(code_digest) != 64: + raise AlgorithmConfigurationError( + "Torch implementation code digest is required" + ) + return TorchStageRunIdentity( + run_id=run_id, + invocation_id=invocation_id, + stage_id=stage_id, + torch_runtime_api_version=1, + algorithm=plan.resolution.algorithm, + implementation_id=plan.implementation.implementation_id, + implementation_code_digest=code_digest, + policy_digest=_policy(plan).digest, + execution_plan_digest=_policy(plan).execution_plan.digest, + plan_digest=plan.plan_id, + ) + + +def _stage_context( + plan: Any, + runtime_context: TorchRuntimeContext, + stage: Any, + index: int, + *, + predecessor: str | None = None, + predecessor_descriptor: Mapping[str, Any] | None = None, +) -> TorchStageContext: + return TorchStageContext( + runtime=runtime_context, + stage_id=stage.stage_id, + stage_index=index, + is_final=stage.stage_id == _policy(plan).execution_plan.final_stage_id, + input_roles=tuple(stage.input_roles), + predecessor_stage_id=predecessor, + predecessor_checkpoint_descriptor=predecessor_descriptor, + metric_mapping=dict(getattr(stage, "metric_mapping", {})), + checkpoint_required=bool(getattr(stage, "checkpoint_required", True)), + checkpoint_interval_windows=int( + getattr(stage, "checkpoint_interval_windows", 1) + ), + ) + + +def _control_for_stage( + plan: Any, + policy: Any, + stage: Any, + *, + run_id: str, + invocation_id: str, + checkpoint: Mapping[str, Any] | None = None, + purpose: str | None = None, + source_stage_id: str | None = None, + predecessor: Mapping[str, Any] | None = None, +) -> dict[str, Any] | None: + """Create credential-free initial recovery control for a Stage.""" + if checkpoint is None: + checkpoint = predecessor + if checkpoint is not None and purpose is None: + purpose = "stage_dependency" + if checkpoint is not None and source_stage_id is None: + source_stage_id = getattr(stage, "checkpoint_from_stage", None) + if checkpoint is None: + return None + resume_uri = checkpoint.get("locator") + descriptor_digest = checkpoint.get("descriptor_digest") + if not isinstance(resume_uri, str) or not resume_uri: + raise AlgorithmConfigurationError( + "Torch recovery requires a credential-free locator" + ) + if not isinstance(descriptor_digest, str): + raise AlgorithmConfigurationError( + "Torch recovery requires checkpoint_descriptor_digest" + ) + locator = TorchCheckpointLocator(resume_uri, descriptor_digest) + control = TorchWorkerControlEnvelope( + schema_version=1, + run_id=run_id, + invocation_id=invocation_id, + source_stage_id=source_stage_id, + target_stage_id=stage.stage_id, + purpose=purpose or "cross_run_initial_recovery", + checkpoint_locator=locator, + checkpoint_descriptor_digest=descriptor_digest, + policy_digest=policy.digest, + execution_plan_digest=policy.execution_plan.digest, + ) + return control.to_dict() + + +def _describe_recovery_locator( + locator: TorchCheckpointLocator, + *, + policy: Any, + plan: Any, + worker_count: int, +) -> TorchCheckpointDescriptor: + """Open a recovery locator on the driver and validate its payload digest.""" + checkpoint = open_torch_checkpoint_locator(locator) + try: + descriptor = describe_torch_checkpoint(TorchCheckpointRef(checkpoint), object()) + _require_checkpoint_commit(checkpoint, descriptor) + finally: + closer = getattr(checkpoint, "close", None) + if callable(closer): + closer() + if descriptor.digest != locator.descriptor_digest: + raise AlgorithmExecutionError( + "Torch recovery locator descriptor digest drifted" + ) + if ( + descriptor.policy_digest != policy.digest + or descriptor.execution_plan_digest != policy.execution_plan.digest + or descriptor.world_size != worker_count + or descriptor.implementation_code_digest != plan.implementation.code_digest + or descriptor.identity.plan_digest != plan.plan_id + or descriptor.input_binding_digest != _input_binding_digest(plan) + or descriptor.state_layout != policy.state_layout + ): + raise AlgorithmExecutionError("Torch recovery descriptor identity mismatch") + return descriptor + + +def _checkpoint_evidence_payload(checkpoint: object) -> dict[str, Any]: + """Read the Core-owned, credential-free evidence sidecar when present.""" + opener = getattr(checkpoint, "as_directory", None) + if not callable(opener): + return {} + try: + with opener() as directory: + path = Path(directory) / "torch_execution_evidence.json" + if path.is_symlink() or not path.is_file(): + return {} + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, TypeError, ValueError) as exc: + raise AlgorithmExecutionError( + "Torch checkpoint execution evidence is malformed" + ) from exc + if not isinstance(payload, Mapping): + raise AlgorithmExecutionError( + "Torch checkpoint execution evidence is malformed" + ) + return dict(payload) + + +def _require_checkpoint_commit( + checkpoint: object, descriptor: TorchCheckpointDescriptor +) -> None: + """Require the marker-last commit used by persistent Stage locators.""" + with _opened_checkpoint(checkpoint) as root: + marker = root / "torch_stage_commit.json" + if marker.is_symlink() or not marker.is_file(): + raise AlgorithmExecutionError( + "Torch Stage locator references an uncommitted checkpoint" + ) + try: + payload = json.loads(marker.read_text(encoding="utf-8")) + except (OSError, TypeError, ValueError) as exc: + raise AlgorithmExecutionError( + "Torch Stage checkpoint commit marker is malformed" + ) from exc + if not isinstance(payload, Mapping) or ( + payload.get("identity") != descriptor.identity.to_dict() + or payload.get("descriptor_digest") != descriptor.digest + ): + raise AlgorithmExecutionError( + "Torch Stage checkpoint commit marker does not match descriptor" + ) + + +def _recovery_record_for_locator( + locator: TorchCheckpointLocator, + *, + policy: Any, + plan: Any, + worker_count: int, +) -> dict[str, Any]: + descriptor = _describe_recovery_locator( + locator, policy=policy, plan=plan, worker_count=worker_count + ) + checkpoint = open_torch_checkpoint_locator(locator) + try: + evidence = _checkpoint_evidence_payload(checkpoint) + finally: + closer = getattr(checkpoint, "close", None) + if callable(closer): + closer() + return { + "locator": locator.uri, + "descriptor_digest": descriptor.digest, + "descriptor": descriptor.to_dict(), + "evidence": evidence, + } + + +def _recovery_records( + plan: Any, + policy: Any, + *, + worker_count: int, +) -> tuple[tuple[str, ...], str | None, dict[str, dict[str, Any]]]: + """Normalize the full Torch recovery envelope and legacy shorthand.""" + stages = tuple(policy.execution_plan.stages) + stage_ids = tuple(stage.stage_id for stage in stages) + raw_envelope = plan.runtime.torch_recovery + if raw_envelope is not None: + envelope = TorchRecoveryEnvelope.from_dict(raw_envelope) + if not policy.resume_supported and ( + envelope.stage_checkpoints or envelope.active_checkpoint is not None + ): + raise AlgorithmConfigurationError( + "Torch Policy does not support external recovery" + ) + completed = tuple(envelope.completed_stage_ids) + if any(stage_id not in stage_ids for stage_id in completed): + raise AlgorithmExecutionError( + "Torch recovery contains an unknown completed Stage" + ) + if tuple(stage_ids[: len(completed)]) != completed: + raise AlgorithmExecutionError( + "Torch recovery completed stages must follow execution plan order" + ) + if ( + envelope.active_stage_id is not None + and envelope.active_stage_id not in stage_ids + ): + raise AlgorithmExecutionError("Torch recovery active Stage is unknown") + if envelope.active_stage_id is not None: + active_index = stage_ids.index(envelope.active_stage_id) + if any(stage_id not in completed for stage_id in stage_ids[:active_index]): + raise AlgorithmExecutionError( + "Torch recovery active Stage has an incomplete predecessor" + ) + active_stage = stages[active_index] + if any(dep not in completed for dep in active_stage.depends_on): + raise AlgorithmExecutionError( + "Torch recovery active Stage dependencies are incomplete" + ) + records: dict[str, dict[str, Any]] = {} + for stage_id, locator in envelope.stage_checkpoints.items(): + record = _recovery_record_for_locator( + locator, policy=policy, plan=plan, worker_count=worker_count + ) + descriptor = TorchCheckpointDescriptor.from_dict(record["descriptor"]) + if descriptor.identity.stage_id != stage_id: + raise AlgorithmExecutionError( + "Torch recovery checkpoint Stage mismatch" + ) + records[stage_id] = record + if ( + envelope.active_stage_id is not None + and envelope.active_checkpoint is not None + ): + active_record = _recovery_record_for_locator( + envelope.active_checkpoint, + policy=policy, + plan=plan, + worker_count=worker_count, + ) + active_descriptor = TorchCheckpointDescriptor.from_dict( + active_record["descriptor"] + ) + if active_descriptor.identity.stage_id != envelope.active_stage_id: + raise AlgorithmExecutionError( + "Torch recovery active checkpoint Stage mismatch" + ) + if not policy.resume_supported: + raise AlgorithmConfigurationError( + "Torch Policy does not support cross-Run active recovery" + ) + if not active_descriptor.resume_supported: + raise AlgorithmExecutionError( + "Torch active recovery checkpoint is not externally recoverable" + ) + records[envelope.active_stage_id] = active_record + return completed, envelope.active_stage_id, records + + resume_uri = plan.runtime.resume_from + ray_config = plan.algorithm_config.get("ray", {}) + resume_config = ( + ray_config.get("resume", {}) if isinstance(ray_config, Mapping) else {} + ) + if resume_uri is None and isinstance(resume_config, Mapping): + for key in ("uri", "checkpoint_uri"): + if isinstance(resume_config.get(key), str): + resume_uri = resume_config[key] + break + if resume_uri is None: + return (), None, {} + if not policy.resume_supported: + raise AlgorithmConfigurationError( + "Torch Policy does not support external recovery" + ) + descriptor_digest = ( + resume_config.get("checkpoint_descriptor_digest") + if isinstance(resume_config, Mapping) + else None + ) + if not isinstance(descriptor_digest, str): + raise AlgorithmConfigurationError( + "Torch resume shorthand requires ray.resume.checkpoint_descriptor_digest" + ) + locator = TorchCheckpointLocator(resume_uri, descriptor_digest) + record = _recovery_record_for_locator( + locator, policy=policy, plan=plan, worker_count=worker_count + ) + descriptor = TorchCheckpointDescriptor.from_dict(record["descriptor"]) + if descriptor.identity.stage_id not in stage_ids: + raise AlgorithmExecutionError( + "Torch resume checkpoint Stage is not in execution plan" + ) + if not descriptor.resume_supported: + raise AlgorithmExecutionError( + "Torch resume checkpoint is not externally recoverable" + ) + return (), descriptor.identity.stage_id, {descriptor.identity.stage_id: record} + + +def _payload_rows(value: object) -> int | None: + count = getattr(value, "count", None) + if callable(count): + try: + result = count() + return int(result) if result is not None else None + except Exception as exc: + logger.debug("Torch Dataset row-count probe failed: %s", type(exc).__name__) + return None + return None + + +def _prepare_datasets(envelope: RuntimeExecutionEnvelope) -> PreparedInput: + adapter = _load_reference(envelope.plan.runtime.worker_input_adapter_ref) + if not callable(adapter): + raise AlgorithmConfigurationError("Torch Worker input adapter is not callable") + prepared = adapter(envelope.input_payloads[0]) + if not isinstance(prepared, PreparedInput) or not prepared.views: + if isinstance(prepared, PreparedInput): + prepared.close() + raise AlgorithmConfigurationError( + "Torch input adapter did not expose Ray datasets" + ) + return prepared + + +def _resource_map(plan: Any) -> dict[str, float]: + resources: dict[str, float] = {"CPU": plan.runtime.num_cpus} + if plan.runtime.num_gpus: + resources["GPU"] = plan.runtime.num_gpus + if plan.runtime.memory_bytes is not None: + resources["memory"] = plan.runtime.memory_bytes + resources.update(plan.runtime.custom_resources) + return resources + + +@contextmanager +def _opened_checkpoint(checkpoint: object) -> Generator[Path, None, None]: + if isinstance(checkpoint, (str, Path)): + root = Path(checkpoint) + if not root.is_dir(): + raise AlgorithmExecutionError("Torch checkpoint directory is missing") + yield root + return + opener = getattr(checkpoint, "as_directory", None) + if not callable(opener): + raise AlgorithmExecutionError("Torch checkpoint cannot be opened") + with opener() as directory: + yield Path(directory) + + +def _validate_stage_routes( + policy: Any, + stage: Any, + datasets: Mapping[str, object], + worker_count: int, +) -> dict[str, int]: + """Validate role presence, exact coverage minimums and replication budgets.""" + routes = {route.role: route for route in policy.dataset_routing} + rows: dict[str, int] = {} + replicated_bytes = 0 + for role in stage.input_roles: + route = routes.get(role) + if route is None: + raise AlgorithmConfigurationError(f"Torch Stage role {role!r} has no route") + dataset = datasets.get(role) + if dataset is None: + if route.required: + raise AlgorithmConfigurationError( + f"required Torch role {role!r} is absent" + ) + continue + count = _payload_rows(dataset) + if count is None: + raise AlgorithmConfigurationError( + f"Torch role {role!r} row count is not verifiable" + ) + rows[role] = count + if count < route.min_total_rows_if_present: + raise AlgorithmConfigurationError(f"Torch role {role!r} has too few rows") + if ( + route.mode == "split_exact" + and route.required + and count < worker_count * route.min_rows_per_worker + ): + raise AlgorithmConfigurationError( + f"Torch role {role!r} cannot cover every worker" + ) + if route.mode == "replicate": + if route.max_rows is None or count > route.max_rows: + raise AlgorithmConfigurationError( + f"Torch replicate role {role!r} exceeds max_rows" + ) + limiter = getattr(dataset, "limit", None) + if not callable(limiter): + raise AlgorithmConfigurationError( + f"Torch replicate role {role!r} cannot be bounded" + ) + probe = limiter(route.max_rows + 1) + probe_count = getattr(probe, "count", None) + if not callable(probe_count): + raise AlgorithmConfigurationError( + f"Torch replicate role {role!r} bounded size is not verifiable" + ) + observed_probe = int(probe_count()) + if observed_probe > route.max_rows: + raise AlgorithmConfigurationError( + f"Torch replicate role {role!r} exceeds max_rows" + ) + if isinstance(datasets, dict): + datasets[role] = limiter(route.max_rows) + size_bytes = getattr(dataset, "size_bytes", None) + if callable(size_bytes): + size_bytes = size_bytes() + if not isinstance(size_bytes, int) or size_bytes < 0: + raise AlgorithmConfigurationError( + f"Torch replicate role {role!r} size is not verifiable" + ) + if ( + route.max_bytes_per_worker is None + or size_bytes > route.max_bytes_per_worker + ): + raise AlgorithmConfigurationError( + f"Torch replicate role {role!r} exceeds max_bytes_per_worker" + ) + replicated_bytes += size_bytes + if ( + policy.max_replicated_bytes_per_worker is not None + and replicated_bytes > policy.max_replicated_bytes_per_worker + ): + raise AlgorithmConfigurationError( + "Torch replicate roles exceed aggregate byte budget" + ) + return rows + + +def _worker_rows(payload: object, role: str) -> int: + if hasattr(payload, "get") and callable(payload.get): + payload = payload.get(role) + value = getattr(payload, "value", payload) + rows = _payload_rows(value) + return rows if rows is not None else 0 + + +def _worker_evidence( + metrics: Mapping[str, Any], + plan: Any, + identity: TorchStageRunIdentity, + stage: Any, + expected_rows: Mapping[str, int], +) -> tuple[dict[str, Any], ...]: + raw_workers = metrics.get("execution_workers") + if ( + not isinstance(raw_workers, (list, tuple)) + or len(raw_workers) != plan.runtime.worker_count + ): + raise AlgorithmExecutionError("Torch execution did not report every worker") + workers = tuple( + WorkerExecutionEvidence.from_dict(item) + for item in _normalize_worker_evidence(raw_workers, plan) + ) + role_evidence: list[TorchRoleExecutionEvidence] = [] + for role in stage.input_roles: + observed = sum(item.input_rows.get(role, 0) for item in workers) + role_evidence.append( + TorchRoleExecutionEvidence( + role=role, + mode="split_exact", + required=True, + present=True, + empty_rank_policy="reject", + expected_rows=expected_rows.get(role), + observed_rows=observed, + rows_per_rank=tuple(item.input_rows.get(role, 0) for item in workers), + binding_digest=_binding_digest_for_role(plan, role), + ) + ) + global_digest = metrics.get("model_state_digest") + if not isinstance(global_digest, str) or len(global_digest) != 64: + raise AlgorithmExecutionError("Torch execution did not report a model digest") + evidence = TorchExecutionEvidence( + identity=identity, + run_config_name=torch_run_config_name(identity), + policy_digest=_policy(plan).digest, + parallelism_id=_policy(plan).parallelism_id, + state_layout=_policy(plan).state_layout, + workers=workers, + roles=tuple(role_evidence), + replicated_state=( + __import__( + "tributo.algorithms.api.execution", + fromlist=["ReplicatedTorchStateEvidence"], + ).ReplicatedTorchStateEvidence( + model_digests_by_rank={ + item.rank: item.model_state_digest for item in workers + }, + global_model_digest=global_digest, + ) + if _policy(plan).state_layout == "replicated" + else None + ), + ) + return (evidence.to_dict(),) + + +def _binding_digest_for_role(plan: Any, role: str) -> str: + """Resolve a role digest, falling back to the primary binding for aliases.""" + if hasattr(plan.input_descriptors, "get"): + try: + descriptor = plan.input_descriptors.get(role) + except AlgorithmConfigurationError: + descriptor = None + if descriptor is not None: + return cast(str, descriptor.binding_digest) + return cast(str, plan.primary_input_descriptor.binding_digest) + + +def _input_binding_digest(plan: Any) -> str: + """Digest the complete role-keyed input descriptor set.""" + if not hasattr(plan, "input_descriptors"): + bindings = cast(tuple[Any, ...], getattr(plan.input_bindings, "bindings", ())) + if len(bindings) == 1: + return hashlib.sha256( + json.dumps( + bindings[0].descriptor_payload(), + sort_keys=True, + separators=(",", ":"), + ).encode() + ).hexdigest() + raise AlgorithmConfigurationError("Torch input binding descriptors are missing") + descriptors = plan.input_descriptors.to_dict() + if len(descriptors) == 1: + return cast(str, plan.primary_input_descriptor.binding_digest) + return hashlib.sha256( + json.dumps(descriptors, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +def _normalize_worker_evidence( + raw_workers: list[object] | tuple[object, ...], plan: Any +) -> tuple[Mapping[str, object], ...]: + """Add Core-owned declared resources when an Adapter omits that field.""" + resources: dict[str, object] = { + "num_cpus": plan.runtime.num_cpus, + "num_gpus": plan.runtime.num_gpus, + "custom": dict(plan.runtime.custom_resources), + } + if plan.runtime.memory_bytes is not None: + resources["memory_bytes"] = plan.runtime.memory_bytes + normalized: list[Mapping[str, object]] = [] + for item in raw_workers: + if not isinstance(item, Mapping): + raise AlgorithmExecutionError("Torch worker evidence is malformed") + payload = dict(item) + payload.setdefault("resources", resources) + normalized.append(payload) + return tuple(normalized) + + +def _component_stage_evidence( + *, + plan: Any, + policy: Any, + stage: Any, + identity: TorchStageRunIdentity, + metrics: Mapping[str, Any], + expected_rows: Mapping[str, int], +) -> ComponentStageEvidence: + raw_workers = metrics.get("execution_workers") + if not isinstance(raw_workers, (list, tuple)): + raise AlgorithmExecutionError("Torch Stage did not report worker evidence") + workers = tuple( + WorkerExecutionEvidence.from_dict(item) + for item in _normalize_worker_evidence(raw_workers, plan) + ) + if len(workers) != plan.runtime.worker_count: + raise AlgorithmExecutionError("Torch Stage worker evidence is incomplete") + roles = _role_execution_evidence( + plan=plan, + policy=policy, + stage=stage, + workers=workers, + expected_rows=expected_rows, + ) + state_digest = metrics.get("model_state_digest") + descriptor = metrics.get("checkpoint_descriptor") + if not isinstance(state_digest, str) or len(state_digest) != 64: + raise AlgorithmExecutionError("Torch Stage did not report a state digest") + if getattr(stage, "checkpoint_required", True) and not isinstance( + descriptor, Mapping + ): + raise AlgorithmExecutionError( + "Torch Stage did not report a checkpoint descriptor" + ) + checkpoint_descriptor_digest = ( + TorchCheckpointDescriptor.from_dict(descriptor).digest + if isinstance(descriptor, Mapping) + else None + ) + return ComponentStageEvidence( + stage_id=identity.stage_id, + workers=workers, + roles=roles, + state_digest=state_digest, + checkpoint_descriptor_digest=checkpoint_descriptor_digest, + ) + + +def _recovered_stage_evidence( + *, + plan: Any, + policy: Any, + stage: Any, + descriptor: Mapping[str, Any], + evidence: Mapping[str, Any], +) -> ComponentStageEvidence: + """Rebuild component evidence persisted beside a completed Stage checkpoint.""" + workers_value = evidence.get("execution_workers") + state_digest = evidence.get("model_state_digest") + if not isinstance(workers_value, (list, tuple)) or not isinstance( + state_digest, str + ): + raise AlgorithmExecutionError( + f"Torch recovery checkpoint for Stage {stage.stage_id!r} has no evidence" + ) + workers = tuple( + WorkerExecutionEvidence.from_dict(item) + for item in _normalize_worker_evidence(workers_value, plan) + ) + if len(workers) != plan.runtime.worker_count: + raise AlgorithmExecutionError( + "Torch recovery Stage worker evidence is incomplete" + ) + expected_rows: dict[str, int] = {} + routes = {route.role: route for route in policy.dataset_routing} + for role in stage.input_roles: + route = routes[role] + rows = tuple(worker.input_rows.get(role, 0) for worker in workers) + if route.mode == "replicate": + if rows and len(set(rows)) == 1: + expected_rows[role] = rows[0] + elif sum(rows) > 0: + expected_rows[role] = sum(rows) + metrics = dict(evidence) + metrics["checkpoint_descriptor"] = dict(descriptor) + identity = TorchStageRunIdentity.from_dict(descriptor["identity"]) + return _component_stage_evidence( + plan=plan, + policy=policy, + stage=stage, + identity=identity, + metrics=metrics, + expected_rows=expected_rows, + ) + + +def _component_state_details( + stages: tuple[ComponentStageEvidence, ...], +) -> dict[str, str | int]: + """Project component evidence into the scalar Core state receipt details.""" + if not stages: + raise AlgorithmExecutionError("Torch component state requires Stage evidence") + composition_digest = hashlib.sha256( + json.dumps( + [stage.to_dict() for stage in stages], + sort_keys=True, + separators=(",", ":"), + ).encode() + ).hexdigest() + details: dict[str, str | int] = { + "framework": "torch_component", + "component_stage_count": len(stages), + "component_stages": ",".join(stage.stage_id for stage in stages), + "anchor_stage": stages[-1].stage_id, + "composition_digest": composition_digest, + } + for stage in stages: + train_rows = sum( + role.observed_rows + for role in stage.roles + if role.role == "train" and role.present + ) + if train_rows == 0: + train_rows = sum(worker.rows_processed or 0 for worker in stage.workers) + details[f"stage.{stage.stage_id}.digest"] = stage.state_digest + details[f"stage.{stage.stage_id}.rows"] = train_rows + return details + + +def _role_execution_evidence( + *, + plan: Any, + policy: Any, + stage: Any, + workers: tuple[WorkerExecutionEvidence, ...], + expected_rows: Mapping[str, int], +) -> tuple[TorchRoleExecutionEvidence, ...]: + """Build role evidence from the Policy route instead of a default split.""" + routes = {route.role: route for route in policy.dataset_routing} + evidence: list[TorchRoleExecutionEvidence] = [] + for role in stage.input_roles: + route = routes[role] + present = role in expected_rows + rows_per_rank = tuple( + worker.input_rows.get(role, 0) if present else 0 for worker in workers + ) + if route.mode == "replicate" and present: + if len(set(rows_per_rank)) != 1: + raise AlgorithmExecutionError( + f"Torch replicate role {role!r} is not identical across workers" + ) + observed_rows = rows_per_rank[0] + replicated_bytes = route.max_bytes_per_worker + else: + observed_rows = sum(rows_per_rank) + replicated_bytes = None + if ( + present + and route.mode in {"split_exact", "replicate"} + and expected_rows.get(role) != observed_rows + ): + raise AlgorithmExecutionError( + f"Torch role {role!r} worker evidence does not prove exact coverage" + ) + evidence.append( + TorchRoleExecutionEvidence( + role=role, + mode=route.mode, + required=route.required, + present=present, + empty_rank_policy=route.empty_rank_policy, + expected_rows=expected_rows.get(role), + observed_rows=observed_rows, + rows_per_rank=rows_per_rank, + replicated_bytes_per_worker=replicated_bytes, + binding_digest=_binding_digest_for_role(plan, role), + ) + ) + return tuple(evidence) + + +def _reduce_composite_loss( + loss: TorchCompositeLossContribution, + *, + config: Mapping[str, Any], + world_size: int, + device: object, + dist: Any, + observation: dict[str, object] | None = None, + expected_metrics: frozenset[str] = frozenset(), +) -> TorchGlobalLossReduction: + """AllReduce generic component state, then invoke the Wheel-owned reducer.""" + import torch + + reducer_ref = config.get("_core_global_loss_reducer_ref") + schema_id = config.get("_core_composite_loss_schema_id") + if not isinstance(reducer_ref, str) or not isinstance(schema_id, str): + raise AlgorithmExecutionError("Composite loss requires a qualified reducer") + reducer_reference = QualifiedReference.parse(reducer_ref) + expected_code_digest = config.get("_core_global_loss_reducer_code_digest") + _validate_module_digest(reducer_reference, expected_code_digest) + reducer = _load_reference(reducer_reference) + if isinstance(reducer, type): + reducer = reducer() + if ( + getattr(reducer, "api_version", None) + != config.get("_core_global_loss_reducer_api_version") + or getattr(reducer, "component_schema_id", None) != schema_id + or getattr(reducer, "code_digest", None) != expected_code_digest + ): + raise AlgorithmExecutionError("Torch global loss reducer identity drifted") + # All ranks must agree on both component key sets before entering any + # value collective. Otherwise one rank can issue a different number of + # all-reduces and deadlock the entire Stage. + local_component_keys = sorted(str(name) for name in loss.differentiable_components) + local_normalizer_keys = sorted(str(name) for name in loss.normalizer_components) + if dist.is_available() and dist.is_initialized(): + gathered_keys: list[object] = [None] * world_size + dist.all_gather_object( + gathered_keys, (local_component_keys, local_normalizer_keys) + ) + if any( + not isinstance(value, (list, tuple)) + or len(value) != 2 + or list(value[0]) != local_component_keys + or list(value[1]) != local_normalizer_keys + for value in gathered_keys + ): + raise AlgorithmExecutionError( + "Composite loss component keys differ across ranks" + ) + local_normalizers = dict(loss.normalizer_components) + global_components: dict[str, float] = {} + for name, value in loss.differentiable_components.items(): + tensor = cast(Any, value).detach().to(dtype=torch.float64) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + global_components[name] = float(tensor.item()) + tensors: dict[str, torch.Tensor] = {} + for name, value in local_normalizers.items(): + tensor = torch.tensor(float(value), dtype=torch.float64, device=device) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + tensors[name] = tensor + global_state = TorchCompositeGlobalState( + components=dict(global_components), + normalizers={name: float(tensor.item()) for name, tensor in tensors.items()}, + ) + global_normalizer = sum(global_state.normalizers.values()) + if not math.isfinite(global_normalizer) or global_normalizer <= 0: + raise AlgorithmExecutionError( + "Composite loss global normalizer must be positive before backward" + ) + algorithm_config = { + str(key): value + for key, value in config.items() + if not str(key).startswith("_core_") + and str(key) not in {"core_control", "ray", "output"} + } + context = TorchGlobalLossContext( + world_size=world_size, + policy_digest=str(config["_core_policy_digest"]), + execution_plan_digest=str(config["_core_execution_plan_digest"]), + config=algorithm_config, + ) + reduction = invoke_torch_global_loss_reducer( + loss, + global_state, + cast(TorchGlobalLossReducer, reducer), + context, + ) + if dist.is_available() and dist.is_initialized(): + payload = ( + reduction.to_dict() + if hasattr(reduction, "to_dict") + else { + "status": reduction.status, + "coefficients": dict(reduction.coefficients), + "branch": reduction.branch, + "failure_code": reduction.failure_code, + } + ) + gathered: list[object] = [None] * world_size + dist.all_gather_object(gathered, payload) + if any(value != gathered[0] for value in gathered[1:]): + raise AlgorithmExecutionError( + "Composite reducer results differ across ranks" + ) + if reduction.status == "rejected": + raise AlgorithmExecutionError( + f"Composite loss reducer rejected contribution: {reduction.failure_code}" + ) + if not expected_metrics.issubset(set(reduction.metrics)): + raise AlgorithmExecutionError( + "Composite reducer did not return every declared metric" + ) + if "train_loss" not in reduction.metrics: + raise AlgorithmExecutionError( + "Composite reducer must return the required train_loss metric" + ) + if observation is not None: + observation["branch"] = reduction.branch + observation["evidence"] = dict(reduction.evidence) + return reduction + + +def _composite_backward( + loss: TorchCompositeLossContribution, + *, + config: Mapping[str, Any], + world_size: int, + device: object, + dist: Any, + observation: dict[str, object] | None = None, + metric_totals: dict[str, list[float]] | None = None, + expected_metrics: frozenset[str] = frozenset(), +) -> object: + """Reduce a Composite loss and return its Core-scaled backward scalar.""" + reduction = _reduce_composite_loss( + loss, + config=config, + world_size=world_size, + device=device, + dist=dist, + observation=observation, + expected_metrics=expected_metrics, + ) + if metric_totals is not None: + _accumulate_metric_totals(metric_totals, reduction.metrics) + return ( + sum( + reduction.coefficients[name] + * cast(Any, loss.differentiable_components[name]) + for name in loss.differentiable_components + ) + * world_size + ) + + +def _restore_torch_retry_checkpoint( + checkpoint: object, + *, + stage_context: TorchStageContext, + model: object, + optimizer: object, + scheduler: object | None, + scaler: object, + rank: int, + strict_identity: bool = True, + progress_sink: dict[str, object] | None = None, + expected_accumulation: int | None = None, +) -> int: + """Validate a Ray-injected retry checkpoint before loading any state.""" + if checkpoint is None: + return 0 + identity = stage_context.runtime.run_identity + if identity is None: + raise AlgorithmExecutionError("Torch retry Stage context has no identity") + ref = TorchCheckpointRef(checkpoint) + descriptor = describe_torch_checkpoint( + ref, + object() + if not strict_identity + else TorchCheckpointContext( + stage=stage_context, + run_id=identity.run_id, + invocation_id=identity.invocation_id, + checkpoint_owner="core", + ), + ) + if strict_identity: + validate_torch_retry_identity( + descriptor, + identity, + world_size=stage_context.runtime.world_size, + ) + elif descriptor.world_size != stage_context.runtime.world_size: + raise AlgorithmExecutionError( + "Torch source checkpoint world size does not match current Stage" + ) + with _opened_checkpoint(checkpoint) as root: + import torch + + model_path = root / "model.pt" + optimizer_path = root / "optimizer.pt" + scaler_path = root / "scaler.pt" + rng_path = root / "rng_state.pt" + if any( + path.is_symlink() or not path.is_file() + for path in (model_path, optimizer_path, scaler_path, rng_path) + ): + raise AlgorithmExecutionError( + "Torch retry checkpoint is missing required model/optimizer/scaler/RNG state" + ) + target_model = cast(Any, getattr(model, "module", model)) + target_model.load_state_dict( + torch.load(model_path, map_location="cpu", weights_only=True) + ) + cast(Any, optimizer).load_state_dict( + torch.load(optimizer_path, map_location="cpu", weights_only=True) + ) + cast(Any, scaler).load_state_dict( + torch.load(scaler_path, map_location="cpu", weights_only=True) + ) + scheduler_path = root / "scheduler.pt" + if scheduler is not None: + if scheduler_path.is_symlink() or not scheduler_path.is_file(): + raise AlgorithmExecutionError( + "Torch retry checkpoint is missing configured scheduler state" + ) + cast(Any, scheduler).load_state_dict( + torch.load(scheduler_path, map_location="cpu", weights_only=True) + ) + payload = torch.load(rng_path, map_location="cpu", weights_only=True) + if not isinstance(payload, Mapping) or not isinstance( + payload.get("states"), list + ): + raise AlgorithmExecutionError("Torch retry RNG state is malformed") + states = payload["states"] + if ( + payload.get("world_size") != stage_context.runtime.world_size + or len(states) != stage_context.runtime.world_size + or rank >= len(states) + or not isinstance(states[rank], bytes) + ): + raise AlgorithmExecutionError( + "Torch retry checkpoint is missing the rank RNG state" + ) + rng = torch.frombuffer(states[rank], dtype=torch.uint8).clone() + torch.set_rng_state(rng) + cuda_states = payload.get("cuda_states_by_rank") + if torch.cuda.is_available(): + if ( + not isinstance(cuda_states, list) + or len(cuda_states) != stage_context.runtime.world_size + or rank >= len(cuda_states) + or not isinstance(cuda_states[rank], list) + or any(not isinstance(state, bytes) for state in cuda_states[rank]) + ): + raise AlgorithmExecutionError("Torch retry CUDA RNG state is malformed") + torch.cuda.set_rng_state_all( + [ + torch.frombuffer(state, dtype=torch.uint8).clone() + for state in cuda_states[rank] + ] + ) + progress_path = root / "torch_progress.json" + if progress_path.is_symlink() or not progress_path.is_file(): + raise AlgorithmExecutionError( + "Torch checkpoint is missing deterministic progress state" + ) + try: + progress = TorchCheckpointProgress.from_dict( + json.loads(progress_path.read_text(encoding="utf-8")) + ) + except (OSError, TypeError, ValueError) as exc: + raise AlgorithmExecutionError( + "Torch checkpoint progress state is malformed" + ) from exc + if progress.optimizer_step != descriptor.completed_step: + raise AlgorithmExecutionError( + "Torch checkpoint progress and descriptor step differ" + ) + if ( + expected_accumulation is not None + and progress.accumulation_steps != expected_accumulation + ): + raise AlgorithmExecutionError( + "Torch checkpoint accumulation configuration differs" + ) + if expected_accumulation is not None and ( + len(progress.dataset_cursor_by_rank) != stage_context.runtime.world_size + or str(rank) not in progress.dataset_cursor_by_rank + ): + raise AlgorithmExecutionError( + "Torch checkpoint dataset cursor does not cover every rank" + ) + if expected_accumulation is not None: + raw_rank_statistics = progress.rank_statistics + if ( + len(raw_rank_statistics) != stage_context.runtime.world_size + or str(rank) not in raw_rank_statistics + ): + raise AlgorithmExecutionError( + "Torch checkpoint statistics do not cover every rank" + ) + if progress_sink is not None: + progress_sink.update(progress.to_dict()) + progress_sink["_typed_progress"] = progress + return descriptor.completed_step + + +def _finalize_torch_window( + *, + scaler: Any, + optimizer: Any, + model: Any, + max_gradient_norm: float | None, + scale: float, +) -> None: + """Apply the Core-owned unscale, normalize, clip, step, and reset order.""" + import torch + + scaler.unscale_(optimizer) + for parameter in model.parameters(): + if parameter.grad is not None: + parameter.grad.mul_(scale) + if max_gradient_norm is not None: + torch.nn.utils.clip_grad_norm_(model.parameters(), max_gradient_norm) + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad() + + +def _should_apply_epoch_scheduler( + *, + restore_same_stage: bool, + epoch: int, + restored_epoch: int, + restored_epoch_scheduler_applied: bool, +) -> bool: + """Return whether the epoch boundary still owes one scheduler step.""" + return not ( + restore_same_stage + and epoch == restored_epoch + and restored_epoch_scheduler_applied + ) + + +def _select_worker_checkpoint( + config: Mapping[str, Any], + stage_context: TorchStageContext, +) -> TorchWorkerCheckpointContext: + """Select retry first, then Core control, never the reverse.""" + import ray.train + + retry = ray.train.get_checkpoint() + identity = stage_context.runtime.run_identity + if identity is None: + raise AlgorithmExecutionError("Torch Worker Stage context has no identity") + if retry is not None: + descriptor = describe_torch_checkpoint( + TorchCheckpointRef(retry), + TorchCheckpointContext( + stage=stage_context, + run_id=identity.run_id, + invocation_id=identity.invocation_id, + checkpoint_owner="core", + ), + ) + validate_torch_retry_identity( + descriptor, + identity, + world_size=stage_context.runtime.world_size, + ) + return TorchWorkerCheckpointContext( + stage=stage_context, + source="ray_failure_retry", + checkpoint=TorchCheckpointRef( + retry, + descriptor_digest=descriptor.digest, + source_stage_id=descriptor.identity.stage_id, + descriptor=descriptor, + ), + ) + control_value = config.get("core_control") + if control_value is None: + return TorchWorkerCheckpointContext(stage=stage_context, source="none") + if not isinstance(control_value, Mapping): + raise AlgorithmExecutionError("Torch Core control envelope is malformed") + control = TorchWorkerControlEnvelope.from_dict(control_value) + if ( + control.target_stage_id != stage_context.stage_id + or control.run_id != identity.run_id + or control.invocation_id != identity.invocation_id + or control.policy_digest != stage_context.runtime.policy_digest + or control.execution_plan_digest != stage_context.runtime.execution_plan_digest + ): + raise AlgorithmExecutionError("Torch Core control envelope identity mismatch") + opener = config.get("_core_checkpoint_opener") + if opener is None: + opener_ref = config.get("_core_checkpoint_opener_ref") + if isinstance(opener_ref, str): + opener = _load_reference(QualifiedReference.parse(opener_ref)) + if not callable(opener): + raise AlgorithmExecutionError( + "Torch Core control envelope has no verified checkpoint opener" + ) + checkpoint = opener(control.checkpoint_locator) + descriptor = describe_torch_checkpoint( + TorchCheckpointRef(checkpoint), + object(), + ) + _require_checkpoint_commit(checkpoint, descriptor) + if descriptor.digest != control.checkpoint_descriptor_digest: + raise AlgorithmExecutionError("Torch Core control descriptor mismatch") + expected_binding = stage_context.runtime.input_binding_digest + if ( + descriptor.policy_digest != stage_context.runtime.policy_digest + or descriptor.execution_plan_digest + != stage_context.runtime.execution_plan_digest + or ( + expected_binding is not None + and descriptor.input_binding_digest != expected_binding + ) + or descriptor.identity.implementation_id + != stage_context.runtime.implementation_id + or descriptor.implementation_code_digest != identity.implementation_code_digest + or descriptor.world_size != stage_context.runtime.world_size + or descriptor.state_layout != stage_context.runtime.state_layout + ): + raise AlgorithmExecutionError("Torch Core control descriptor identity mismatch") + if control.purpose == "stage_dependency" and ( + descriptor.identity.stage_id != control.source_stage_id + or control.source_stage_id != stage_context.predecessor_stage_id + or descriptor.identity.run_id != identity.run_id + or descriptor.identity.invocation_id != identity.invocation_id + ): + raise AlgorithmExecutionError("Torch Core control source Stage mismatch") + if ( + control.purpose == "cross_run_initial_recovery" + and not descriptor.resume_supported + ): + raise AlgorithmExecutionError("Torch checkpoint is not externally recoverable") + return TorchWorkerCheckpointContext( + stage=stage_context, + source=( + "stage_dependency" + if control.purpose == "stage_dependency" + else "cross_run_initial_recovery" + ), + checkpoint=TorchCheckpointRef( + checkpoint, + descriptor_digest=descriptor.digest, + source_stage_id=descriptor.identity.stage_id, + descriptor=descriptor, + ), + ) + + +def open_torch_checkpoint_locator(locator: TorchCheckpointLocator) -> object: + """Core-owned locator opener hook; credentials are resolved by the runtime.""" + if not isinstance(locator, TorchCheckpointLocator): + raise AlgorithmExecutionError("Torch checkpoint locator is invalid") + if locator.uri.startswith("ray://"): + from ray.train import Checkpoint + + path = Path(locator.uri.removeprefix("ray://")) + if path.is_dir(): + try: + return Checkpoint.from_directory(str(path)) + except (OSError, ValueError, TypeError) as exc: + raise AlgorithmExecutionError( + "Torch checkpoint locator could not be opened" + ) from exc + if locator.uri.startswith(("s3://", "gs://", "gcs://", "hdfs://")): + from ray.train import Checkpoint + + try: + return Checkpoint(locator.uri) + except (OSError, ValueError, TypeError) as exc: + raise AlgorithmExecutionError( + "Torch checkpoint locator could not be opened" + ) from exc + raise AlgorithmExecutionError( + "Torch checkpoint locator requires a configured Core storage opener" + ) + + +def _persist_stage_checkpoint( + checkpoint: object, + *, + identity: TorchStageRunIdentity, + storage_path: object, + descriptor_digest: str, +) -> str | None: + """Persist a same-invocation Stage checkpoint without replacing prior data.""" + if not isinstance(storage_path, (str, Path)): + return None + raw_storage = str(storage_path) + if "://" in raw_storage and not raw_storage.startswith("file://"): + from urllib.parse import urlsplit + + opener = getattr(checkpoint, "as_directory", None) + if not callable(opener): + raise AlgorithmExecutionError( + "Torch remote Stage checkpoint cannot be opened for Core persistence" + ) + try: + import pyarrow.fs as pafs + + parsed = urlsplit(raw_storage) + filesystem, prefix = pafs.FileSystem.from_uri(raw_storage) + remote_staging = ( + f"{prefix.rstrip('/')}/{identity.run_config_name}/" + f".stage_checkpoint.staging-{descriptor_digest}" + ) + commit_path = f"{remote_staging}/torch_stage_commit.json" + commit_info = filesystem.get_file_info(commit_path) + commit_payload = { + "schema_version": 1, + "identity": identity.to_dict(), + "descriptor_digest": descriptor_digest, + } + if commit_info.type is pafs.FileType.File: + try: + existing_commit = json.loads( + filesystem.open_input_file(commit_path).read().decode("utf-8") + ) + except (OSError, TypeError, ValueError) as exc: + raise AlgorithmExecutionError( + "Torch remote Stage checkpoint commit marker is malformed" + ) from exc + if existing_commit != commit_payload: + raise AlgorithmExecutionError( + "Torch remote Stage checkpoint identity collision" + ) + return f"{parsed.scheme}://{parsed.netloc}/{remote_staging.lstrip('/')}" + filesystem.create_dir(remote_staging, recursive=True) + with opener() as source: + source_path = Path(source) + for source_file in source_path.rglob("*"): + if source_file.is_symlink(): + raise AlgorithmExecutionError( + "Torch Stage checkpoint payload contains a symlink" + ) + if not source_file.is_file(): + continue + relative = source_file.relative_to(source_path).as_posix() + destination = f"{remote_staging.rstrip('/')}/{relative}" + with filesystem.open_output_stream(destination) as stream: + stream.write(source_file.read_bytes()) + with filesystem.open_output_stream(commit_path) as stream: + stream.write( + (json.dumps(commit_payload, sort_keys=True) + "\n").encode() + ) + if not parsed.scheme or not parsed.netloc: + raise AlgorithmExecutionError("Torch remote storage URI is invalid") + return f"{parsed.scheme}://{parsed.netloc}/{remote_staging.lstrip('/')}" + except AlgorithmExecutionError: + raise + except (ImportError, OSError, TypeError, ValueError) as exc: + raise AlgorithmExecutionError( + "failed to persist Torch Stage checkpoint in Core storage" + ) from exc + root = Path(storage_path) + if not root.is_absolute(): + return None + target = root / identity.run_config_name / "stage_checkpoint" + commit_path_local = target / "torch_stage_commit.json" + commit_payload = { + "schema_version": 1, + "identity": identity.to_dict(), + "descriptor_digest": descriptor_digest, + } + if target.exists(): + if target.is_symlink() or not target.is_dir(): + raise AlgorithmExecutionError( + "Torch Stage checkpoint destination is not a directory" + ) + if not commit_path_local.is_file() or commit_path_local.is_symlink(): + raise AlgorithmExecutionError( + "Torch Stage checkpoint destination is a partial or uncommitted snapshot" + ) + try: + if ( + json.loads(commit_path_local.read_text(encoding="utf-8")) + != commit_payload + ): + raise AlgorithmExecutionError( + "Torch Stage checkpoint identity collision" + ) + except (OSError, TypeError, ValueError) as exc: + raise AlgorithmExecutionError( + "Torch Stage checkpoint commit marker is malformed" + ) from exc + return f"ray://{target}" + opener = getattr(checkpoint, "as_directory", None) + if not callable(opener): + return None + temporary_target = Path( + tempfile.mkdtemp( + prefix=f".{target.name}.staging-{descriptor_digest}-", + dir=target.parent, + ) + ) + with opener() as source: + source_path = Path(source) + for source_file in source_path.rglob("*"): + if source_file.is_symlink(): + raise AlgorithmExecutionError( + "Torch Stage checkpoint payload contains a symlink" + ) + if not source_file.is_file(): + continue + relative_local = source_file.relative_to(source_path) + destination_local = temporary_target / relative_local + destination_local.parent.mkdir(parents=True, exist_ok=True) + destination_local.write_bytes(source_file.read_bytes()) + (temporary_target / "torch_stage_commit.json").write_text( + json.dumps(commit_payload, sort_keys=True) + "\n", encoding="utf-8" + ) + os.replace(temporary_target, target) + return f"ray://{target}" + + +def _recipe_worker(config: Mapping[str, Any]) -> None: + """Run a minimal Core-owned Recipe loop in a Ray Train worker.""" + import ray.train + import torch + from ray.train.torch import prepare_model + + recipe_ref = config.get("_core_implementation_ref") + if not isinstance(recipe_ref, str): + raise AlgorithmConfigurationError( + "Core Worker implementation reference is missing" + ) + recipe = _load_reference(QualifiedReference.parse(recipe_ref)) + if not isinstance(recipe, type) or not issubclass(recipe, TorchRecipe): + raise AlgorithmConfigurationError("Core Worker did not receive a TorchRecipe") + recipe_instance = recipe() + stage_context_value = config.get("_core_stage_context") + if not isinstance(stage_context_value, Mapping): + raise AlgorithmConfigurationError("Core Worker stage context is missing") + stage_context = TorchStageContext.from_dict(stage_context_value) + runtime_context = stage_context.runtime + modules = recipe_instance.build_modules( + TorchBuildContext(runtime=runtime_context, stage=stage_context) + ) + if isinstance(modules, TorchModuleSet): + module_set = modules + elif isinstance(modules, Mapping): + module_set = TorchModuleSet(modules) + else: + raise AlgorithmConfigurationError( + "TorchRecipe.build_modules must return TorchModuleSet" + ) + model = module_set["model"] + if not isinstance(model, torch.nn.Module): + raise AlgorithmConfigurationError("TorchRecipe model must be torch.nn.Module") + optimization = recipe_instance.configure_optimizers( + module_set, + TorchBuildContext(runtime=runtime_context, stage=stage_context), + ) + if not isinstance(optimization, TorchOptimizationPlan): + raise AlgorithmConfigurationError( + "TorchRecipe.configure_optimizers returned invalid plan" + ) + model = cast(Any, prepare_model(model)) + module_set = TorchModuleSet({**module_set.modules, "model": model}) + optimizer = cast(Any, optimization.optimizer) + optimizer.zero_grad() + stage_roles = tuple(stage_context.input_roles) + if not stage_roles: + raise AlgorithmConfigurationError("TorchRecipe Stage has no input roles") + role_shards: dict[str, object] = {} + for role in stage_roles: + try: + shard = ray.train.get_dataset_shard(role) + except KeyError: + shard = None + if shard is not None: + role_shards[role] = shard + if stage_roles[0] not in role_shards: + raise AlgorithmConfigurationError( + f"TorchRecipe requires the {stage_roles[0]!r} dataset" + ) + train = role_shards[stage_roles[0]] + training_roles = tuple(role for role in stage_roles if role not in {"val", "test"}) + multi_role = len(training_roles) > 1 + rank = ray.train.get_context().get_world_rank() + world_size = ray.train.get_context().get_world_size() + training = config.get("training", {}) + if not isinstance(training, Mapping): + raise AlgorithmConfigurationError("Torch training config must be a mapping") + epochs = training.get("epochs", 1) + shuffle = training.get("shuffle", False) + if not isinstance(shuffle, bool): + raise AlgorithmConfigurationError("Torch training shuffle must be boolean") + if shuffle: + raise AlgorithmConfigurationError( + "Torch Runtime v1 requires an unshuffled Ray Dataset for exact recovery" + ) + batch_size = training.get("batch_size", config.get("_core_batch_size", 32)) + if not isinstance(epochs, int) or isinstance(epochs, bool) or epochs < 1: + raise AlgorithmConfigurationError("Torch training epochs must be positive") + if ( + not isinstance(batch_size, int) + or isinstance(batch_size, bool) + or batch_size < 1 + ): + raise AlgorithmConfigurationError("Torch training batch_size must be positive") + amp = bool(training.get("amp", False)) + if amp and not torch.cuda.is_available(): + raise AlgorithmConfigurationError("Torch AMP requires CUDA") + scaler = torch.amp.GradScaler("cuda", enabled=amp) + seed = training.get("seed", 42) + if not isinstance(seed, int) or isinstance(seed, bool): + raise AlgorithmConfigurationError("Torch training seed must be an integer") + torch.manual_seed(seed + rank) + scheduler = optimization.scheduler + accumulation = optimization.gradient_accumulation_steps + checkpoint_context = _select_worker_checkpoint(config, stage_context) + restored_progress: dict[str, object] = {} + try: + loaded_step = _restore_torch_retry_checkpoint( + checkpoint_context.checkpoint.checkpoint + if checkpoint_context.checkpoint is not None + else None, + stage_context=stage_context, + model=model, + optimizer=optimization.optimizer, + scheduler=scheduler, + scaler=scaler, + rank=rank, + strict_identity=checkpoint_context.source == "ray_failure_retry", + progress_sink=restored_progress, + expected_accumulation=( + accumulation + if checkpoint_context.source + in {"ray_failure_retry", "cross_run_initial_recovery"} + else None + ), + ) + restored_step = ( + loaded_step + if checkpoint_context.source + in {"ray_failure_retry", "cross_run_initial_recovery"} + else 0 + ) + finally: + if checkpoint_context.checkpoint is not None: + checkpoint_context.checkpoint.close() + restore_same_stage = checkpoint_context.source in { + "ray_failure_retry", + "cross_run_initial_recovery", + } + typed_progress = restored_progress.get("_typed_progress") + if restore_same_stage and not isinstance(typed_progress, TorchCheckpointProgress): + raise AlgorithmExecutionError("Torch checkpoint progress is not typed") + restored_checkpoint_progress = ( + cast(TorchCheckpointProgress, typed_progress) if restore_same_stage else None + ) + rank_statistics = ( + restored_checkpoint_progress.rank_statistics.get(str(rank)) + if restored_checkpoint_progress is not None + else None + ) + if restore_same_stage and not isinstance( + rank_statistics, TorchRankProgressStatistics + ): + raise AlgorithmExecutionError("Torch checkpoint rank statistics are missing") + + rows = rank_statistics.rows_processed if rank_statistics is not None else 0 + steps = 0 + coverage_totals = ( + dict(rank_statistics.coverage_totals) if rank_statistics is not None else {} + ) + loss_numerator_total = ( + rank_statistics.loss_numerator_total if rank_statistics is not None else 0.0 + ) + loss_normalizer_total = ( + rank_statistics.loss_normalizer_total if rank_statistics is not None else 0.0 + ) + metric_totals = ( + {name: list(pair) for name, pair in rank_statistics.metric_totals.items()} + if rank_statistics is not None + else {} + ) + evaluation_totals = ( + {name: list(pair) for name, pair in rank_statistics.evaluation_totals.items()} + if rank_statistics is not None + else {} + ) + reducer_observation: dict[str, object] = ( + dict(rank_statistics.reducer_observation) if rank_statistics is not None else {} + ) + composite_loss_seen = False + batch_context = TorchBatchContext( + stage=stage_context, + input_roles=stage_roles, + feature_names=tuple( + name + for name in config.get("_core_feature_names", ()) + if isinstance(name, str) + ), + label_name=( + config.get("_core_label_name") + if isinstance(config.get("_core_label_name"), str) + else None + ), + weight_name=( + config.get("_core_weight_name") + if isinstance(config.get("_core_weight_name"), str) + else None + ), + ) + checkpoint_interval = config.get("_core_checkpoint_interval_windows", 1) + if ( + not isinstance(checkpoint_interval, int) + or isinstance(checkpoint_interval, bool) + or checkpoint_interval < 1 + ): + raise AlgorithmConfigurationError( + "Torch checkpoint interval must be a positive integer" + ) + import torch.distributed as dist + + def emit_checkpoint( + completed_step: int, + *, + epoch: int, + micro_batch_cursor: int, + scheduler_step: int, + rows_processed: int, + coverage_totals: Mapping[str, int], + loss_numerator_total: float, + loss_normalizer_total: float, + metric_totals: Mapping[str, list[float]], + evaluation_totals: Mapping[str, list[float]], + reducer_observation: Mapping[str, object], + ) -> None: + """Report an optimizer-boundary checkpoint for Ray failure retry.""" + from ray.train import Checkpoint + + identity = runtime_context.run_identity + if identity is None: + raise AlgorithmExecutionError("Torch Worker stage context has no identity") + checkpoint_dir = Path(tempfile.mkdtemp(prefix="tributo_torch_checkpoint_")) + try: + target_model = getattr(model, "module", model) + torch.save(target_model.state_dict(), checkpoint_dir / "model.pt") + torch.save( + cast(Any, optimizer).state_dict(), checkpoint_dir / "optimizer.pt" + ) + torch.save(cast(Any, scaler).state_dict(), checkpoint_dir / "scaler.pt") + if scheduler is not None: + torch.save( + cast(Any, scheduler).state_dict(), checkpoint_dir / "scheduler.pt" + ) + cursor_by_rank: list[object] = [None] * world_size + if dist.is_available() and dist.is_initialized(): + dist.all_gather_object(cursor_by_rank, micro_batch_cursor) + else: + cursor_by_rank = [micro_batch_cursor] + if any( + not isinstance(value, int) or isinstance(value, bool) or value < 0 + for value in cursor_by_rank + ): + raise AlgorithmExecutionError( + "Torch dataset cursor collective is incomplete" + ) + local_statistics = TorchRankProgressStatistics( + rows_processed=rows_processed, + coverage_totals=coverage_totals, + loss_numerator_total=loss_numerator_total, + loss_normalizer_total=loss_normalizer_total, + metric_totals={ + name: (values[0], values[1]) + for name, values in metric_totals.items() + }, + evaluation_totals={ + name: (values[0], values[1]) + for name, values in evaluation_totals.items() + }, + reducer_observation=reducer_observation, + ) + statistics_by_rank: list[object] = [None] * world_size + if dist.is_available() and dist.is_initialized(): + dist.all_gather_object(statistics_by_rank, local_statistics.to_dict()) + else: + statistics_by_rank = [local_statistics.to_dict()] + if any(not isinstance(value, Mapping) for value in statistics_by_rank): + raise AlgorithmExecutionError( + "Torch checkpoint statistics collective is incomplete" + ) + progress = TorchCheckpointProgress( + epoch=epoch, + micro_batch_cursor=micro_batch_cursor, + optimizer_step=completed_step, + scheduler_step=scheduler_step, + accumulation_steps=accumulation, + dataset_cursor_by_rank={ + str(rank_id): cast(int, cursor) + for rank_id, cursor in enumerate(cursor_by_rank) + }, + shuffle_seed=int(seed + epoch), + rows_processed=rows_processed, + coverage_totals=coverage_totals, + loss_numerator_total=loss_numerator_total, + loss_normalizer_total=loss_normalizer_total, + metric_totals={ + name: (values[0], values[1]) + for name, values in metric_totals.items() + }, + evaluation_totals={ + name: (values[0], values[1]) + for name, values in evaluation_totals.items() + }, + rank_statistics={ + str(rank_id): TorchRankProgressStatistics.from_dict( + cast(Mapping[str, Any], stats) + ) + for rank_id, stats in enumerate(statistics_by_rank) + }, + epoch_scheduler_applied=False, + ) + (checkpoint_dir / "torch_progress.json").write_text( + json.dumps(progress.to_dict(), sort_keys=True, separators=(",", ":")), + encoding="utf-8", + ) + cpu_rng = torch.get_rng_state().cpu().numpy().tobytes() + cpu_states: list[bytes | None] = [None] * world_size + if dist.is_available() and dist.is_initialized(): + dist.all_gather_object(cpu_states, cpu_rng) + else: + cpu_states = [cpu_rng] + if any(not isinstance(value, bytes) for value in cpu_states): + raise AlgorithmExecutionError( + "Torch RNG state collective is incomplete" + ) + cuda_states_by_rank: list[list[bytes]] = [] + if torch.cuda.is_available(): + local_cuda = [ + state.cpu().numpy().tobytes() + for state in torch.cuda.get_rng_state_all() + ] + gathered_cuda: list[object] = [None] * world_size + if dist.is_available() and dist.is_initialized(): + dist.all_gather_object(gathered_cuda, local_cuda) + if any( + not isinstance(value, list) + or any(not isinstance(state, bytes) for state in value) + for value in gathered_cuda + ): + raise AlgorithmExecutionError( + "Torch CUDA RNG state collective is incomplete" + ) + cuda_states_by_rank = cast(list[list[bytes]], gathered_cuda) + else: + cuda_states_by_rank = [local_cuda] + torch.save( + { + "world_size": world_size, + "states": cast(list[bytes], cpu_states), + "cuda_states_by_rank": cuda_states_by_rank, + }, + checkpoint_dir / "rng_state.pt", + ) + payload_names = [ + "model.pt", + "optimizer.pt", + "scaler.pt", + "rng_state.pt", + "torch_progress.json", + ] + if scheduler is not None: + payload_names.append("scheduler.pt") + descriptor = TorchCheckpointDescriptor( + schema_version=1, + identity=identity, + run_config_name=identity.run_config_name, + state_layout=str(config.get("_core_state_layout", "replicated")), + world_size=world_size, + completed_step=completed_step, + policy_digest=runtime_context.policy_digest, + execution_plan_digest=runtime_context.execution_plan_digest, + input_binding_digest=str(config.get("_core_input_binding_digest", "")), + implementation_code_digest=str( + config.get("_core_implementation_code_digest", "") + ), + payload_files={ + name: hashlib.sha256( + (checkpoint_dir / name).read_bytes() + ).hexdigest() + for name in payload_names + }, + adapter_identity=config.get("_core_adapter_identity"), + resume_supported=runtime_context.resume_supported, + same_world_size_resume=runtime_context.same_world_size_resume, + ) + + class _IntervalDraft: + checkpoint_dir: str | os.PathLike[str] + + def __init__(self) -> None: + self.checkpoint_dir = str(checkpoint_dir) + + def report( + self, + *, + metrics: Mapping[str, object], + stage_context: object, + completed_step: int, + ) -> None: + del stage_context, completed_step + checkpoint = ( + Checkpoint.from_directory(str(checkpoint_dir)) + if rank == int(config.get("_core_checkpoint_owner_rank", 0)) + else None + ) + ray.train.report(dict(metrics), checkpoint=checkpoint) + + report_torch_checkpoint( + { + "train_loss": ( + loss_numerator_total / loss_normalizer_total + if loss_normalizer_total > 0 + else 0.0 + ), + "model_state_digest": hashlib.sha256( + json.dumps( + { + name: str( + cast(Any, value).detach().cpu().numpy().tobytes() + ) + for name, value in target_model.state_dict().items() + }, + sort_keys=True, + ).encode() + ).hexdigest(), + "checkpoint_descriptor": descriptor.to_dict(), + }, + _IntervalDraft(), + stage_context, + completed_step, + ) + finally: + import shutil + + shutil.rmtree(checkpoint_dir, ignore_errors=True) + + def reduce_window(value: float) -> float: + tensor = torch.tensor( + value, dtype=torch.float64, device=next(model.parameters()).device + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + return float(tensor.item()) + + def zero_batch(template: TorchBatch) -> TorchBatch: + def zeros(value: object) -> object: + if not isinstance(value, torch.Tensor): + return value + # Preserve the typed batch signature while making every tensor + # genuinely zero-row. This prevents empty ranks from producing + # synthetic loss/metric contributions in Recipe or reducer code. + if value.ndim > 0: + return value[:0] + return torch.zeros_like(value) + + return TorchBatch( + positional=tuple(zeros(value) for value in template.positional), + keyword={name: zeros(value) for name, value in template.keyword.items()}, + targets=zeros(template.targets) if template.targets is not None else None, + weights=zeros(template.weights) if template.weights is not None else None, + local_rows=0, + coverage_counts=dict.fromkeys(template.coverage_counts, 0), + ) + + def aligned_metric_contributions( + contributions: Mapping[str, TorchMetricContribution], + ) -> dict[str, TorchMetricContribution]: + """Make metric keys identical before any later rank-wise reduction.""" + for _name, contribution in contributions.items(): + if not isinstance(contribution, TorchMetricContribution): + raise AlgorithmConfigurationError( + "TorchStepResult metrics must be TorchMetricContribution values" + ) + names = set(contributions) + if dist.is_available() and dist.is_initialized(): + gathered: list[object] = [None] * world_size + dist.all_gather_object(gathered, sorted(names)) + for value in gathered: + if not isinstance(value, list) or any( + not isinstance(name, str) for name in value + ): + raise AlgorithmExecutionError( + "Torch metric key collective is incomplete" + ) + names.update(value) + return { + name: contributions.get(name, TorchMetricContribution(0.0, 0.0)) + for name in sorted(names) + } + + raw_metric_mapping = config.get("_core_metric_mapping", {}) + metric_mapping = ( + {str(key): str(value) for key, value in raw_metric_mapping.items()} + if isinstance(raw_metric_mapping, Mapping) + else {} + ) + if len(set(metric_mapping.values())) != len(metric_mapping): + raise AlgorithmConfigurationError("Torch metric mapping targets must be unique") + raw_metric_reducers = config.get("_core_metric_reducers", {}) + if not isinstance(raw_metric_reducers, Mapping): + raise AlgorithmConfigurationError("Torch metric reducers must be a mapping") + declared_metric_names = {str(name) for name in raw_metric_reducers} - {"train_loss"} + expected_metric_names = frozenset( + { + source + for source, target in metric_mapping.items() + if target in declared_metric_names + } + | { + name + for name in declared_metric_names + if name not in metric_mapping.values() + } + ) + restored_epoch = ( + restored_checkpoint_progress.epoch + if restored_checkpoint_progress is not None + else 0 + ) + remaining_skip_micro_batches = ( + restored_checkpoint_progress.dataset_cursor_by_rank[str(rank)] + if restored_checkpoint_progress is not None + else 0 + ) + scheduler_steps = ( + restored_checkpoint_progress.scheduler_step + if restored_checkpoint_progress is not None + else 0 + ) + restored_epoch_scheduler_applied = ( + restored_checkpoint_progress.epoch_scheduler_applied + if restored_checkpoint_progress is not None + else False + ) + if restore_same_stage and restored_epoch >= epochs: + raise AlgorithmExecutionError( + "Torch checkpoint progress points past the configured epoch range" + ) + if ( + checkpoint_context.source + in { + "ray_failure_retry", + "cross_run_initial_recovery", + } + and restored_progress.get("shuffle_seed") != seed + restored_epoch + ): + raise AlgorithmExecutionError( + "Torch checkpoint shuffle state does not match the current seed" + ) + + def map_metric_contributions( + contributions: Mapping[str, TorchMetricContribution], + ) -> dict[str, TorchMetricContribution]: + return { + metric_mapping.get(name, name): contribution + for name, contribution in contributions.items() + } + + epoch_micro_batch_cursor = 0 + for _epoch in range(epochs): + if _epoch < restored_epoch: + continue + epoch_micro_batch_cursor = ( + remaining_skip_micro_batches if _epoch == restored_epoch else 0 + ) + next_payload: Any + if multi_role: + role_iterators = { + role: iter( + cast(Any, shard).iter_torch_batches( + batch_size=batch_size, drop_last=False + ) + ) + for role, shard in role_shards.items() + if role in training_roles + } + + def _next_multi_payload( + role_iterators: Mapping[str, Any] = role_iterators, + ) -> object | None: + payload = { + role: next(iterator, None) + for role, iterator in role_iterators.items() + } + return ( + payload + if any(value is not None for value in payload.values()) + else None + ) + + next_payload = _next_multi_payload + + else: + iterator = iter( + cast(Any, train).iter_torch_batches( + batch_size=batch_size, drop_last=False + ) + ) + + def _next_single_payload(iterator: Any = iterator) -> object | None: + return next(iterator, None) + + next_payload = _next_single_payload + + raw = next_payload() + first_active = torch.tensor( + 1 if raw is not None else 0, + dtype=torch.int64, + device=next(model.parameters()).device, + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(first_active, op=dist.ReduceOp.SUM) + if int(first_active.item()) == 0: + continue + if remaining_skip_micro_batches == 0 and int(first_active.item()) != world_size: + raise AlgorithmExecutionError( + "TorchRecipe requires at least one batch on every rank" + ) + template = recipe_instance.adapt_batch(raw, batch_context) + if not isinstance(template, TorchBatch): + raise AlgorithmConfigurationError( + "TorchRecipe.adapt_batch must return TorchBatch" + ) + window = TorchAccumulationWindow( + index=restored_step + steps // accumulation, + expected_micro_batches=accumulation, + ) + while True: + active = raw is not None + active_tensor = torch.tensor( + 1 if active else 0, + dtype=torch.int64, + device=next(model.parameters()).device, + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(active_tensor, op=dist.ReduceOp.SUM) + if int(active_tensor.item()) == 0: + break + next_raw = next_payload() + if remaining_skip_micro_batches > 0: + remaining_skip_micro_batches -= 1 + epoch_micro_batch_cursor += 1 + raw = next_raw + continue + next_active = torch.tensor( + 1 if next_raw is not None else 0, + dtype=torch.int64, + device=next(model.parameters()).device, + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(next_active, op=dist.ReduceOp.SUM) + expected_micro_batches = window.expected_micro_batches + if ( + int(next_active.item()) == 0 + and window.observed_micro_batches + 1 < expected_micro_batches + ): + expected_micro_batches = window.observed_micro_batches + 1 + window = TorchAccumulationWindow( + index=window.index, + expected_micro_batches=expected_micro_batches, + observed_micro_batches=window.observed_micro_batches, + normalizer_total=window.normalizer_total, + ) + batch = ( + recipe_instance.adapt_batch(raw, batch_context) + if active + else zero_batch(template) + ) + if not isinstance(batch, TorchBatch): + raise AlgorithmConfigurationError( + "TorchRecipe.adapt_batch must return TorchBatch" + ) + is_boundary = ( + window.observed_micro_batches + 1 == window.expected_micro_batches + ) + sync_context = ( + model.no_sync() + if hasattr(model, "no_sync") and not is_boundary + else nullcontext() + ) + with ( + sync_context, + torch.autocast(device_type="cuda" if amp else "cpu", enabled=amp), + ): + # Every rank invokes the typed step, including a synchronized + # zero-contribution batch after its shard is exhausted. This + # keeps composite reducer collectives aligned across ranks. + step = recipe_instance.training_step( + module_set, + batch, + TorchStepContext( + stage=stage_context, + window_index=restored_step + steps // accumulation, + micro_batch_index=steps % accumulation, + ), + ) + if not isinstance(step, TorchStepResult): + raise AlgorithmConfigurationError( + "TorchRecipe.training_step returned invalid result" + ) + for name, count in step.coverage_counts.items(): + if name in stage_roles: + raise AlgorithmConfigurationError( + "Torch coverage_counts cannot override Core role rows" + ) + coverage_totals[name] = coverage_totals.get(name, 0) + count + is_composite = isinstance(step.loss, TorchCompositeLossContribution) + if is_composite and accumulation != 1: + raise AlgorithmConfigurationError( + "Core TorchRecipe composite loss requires accumulation_steps=1" + ) + if is_composite: + composite_loss_seen = True + if isinstance(step.loss, TorchCompositeLossContribution): + normalizer = sum( + float(value) + for value in step.loss.normalizer_components.values() + ) + elif isinstance(step.loss, TorchLossContribution): + numerator = step.loss.numerator + normalizer = float(step.loss.normalizer) + loss_numerator_total += float(cast(Any, numerator).detach().item()) + loss_normalizer_total += normalizer + else: + raise AlgorithmConfigurationError( + "TorchRecipe loss contribution is invalid" + ) + for metric_name, contribution in aligned_metric_contributions( + map_metric_contributions(step.metrics) + ).items(): + totals = metric_totals.setdefault(metric_name, [0.0, 0.0]) + totals[0] += contribution.numerator + totals[1] += contribution.normalizer + mapped_step_metrics = map_metric_contributions(step.metrics) + if ( + not isinstance(step.loss, TorchCompositeLossContribution) + and set(mapped_step_metrics) != declared_metric_names + ): + raise AlgorithmExecutionError( + "TorchStepResult.metrics do not match TorchMetricPlan" + ) + backward_context = TorchBackwardContext( + world_size=world_size, + backward=lambda value: scaler.scale(value).backward(), + reduce_normalizer=reduce_window, + reduce_window_normalizer=reduce_window, + compose_composite=( + lambda value: _composite_backward( + value, + config=config, + world_size=world_size, + device=next(model.parameters()).device, + dist=dist, + observation=reducer_observation, + metric_totals=metric_totals, + expected_metrics=expected_metric_names, + ) + ), + finalize_window=lambda scale: _finalize_torch_window( + scaler=scaler, + optimizer=optimizer, + model=model, + max_gradient_norm=optimization.max_gradient_norm, + scale=scale, + ), + ) + # The public helper owns normalizer accumulation and the + # optimizer-window boundary; the callback owns only Torch's + # unscale/clip/step sequence. + result = apply_torch_loss_backward( + step.loss, + window, + backward_context, + ) + if result.window_complete: + window = TorchAccumulationWindow( + index=window.index + 1, + expected_micro_batches=accumulation, + ) + else: + window = window.add(result.local_normalizer) + rows += batch.local_rows + steps += 1 + epoch_micro_batch_cursor += 1 + if ( + result.window_complete + and (steps // accumulation) % checkpoint_interval == 0 + ): + emit_checkpoint( + restored_step + steps // accumulation, + epoch=_epoch, + micro_batch_cursor=epoch_micro_batch_cursor, + scheduler_step=scheduler_steps, + rows_processed=rows, + coverage_totals=coverage_totals, + loss_numerator_total=loss_numerator_total, + loss_normalizer_total=loss_normalizer_total, + metric_totals=metric_totals, + evaluation_totals=evaluation_totals, + reducer_observation=reducer_observation, + ) + raw = next_raw + if optimization.scheduler is not None and _should_apply_epoch_scheduler( + restore_same_stage=restore_same_stage, + epoch=_epoch, + restored_epoch=restored_epoch, + restored_epoch_scheduler_applied=restored_epoch_scheduler_applied, + ): + scheduler_step = getattr(optimization.scheduler, "step", None) + if not callable(scheduler_step): + raise AlgorithmConfigurationError( + "Torch scheduler must implement step()" + ) + scheduler_step() + scheduler_steps += 1 + for split in ("val", "test"): + try: + evaluation_data = ray.train.get_dataset_shard(split) + except KeyError: + evaluation_data = None + if evaluation_data is None: + continue + iterator = iter( + evaluation_data.iter_torch_batches(batch_size=batch_size, drop_last=False) + ) + raw = next(iterator, None) + active = torch.tensor( + 1 if raw is not None else 0, + dtype=torch.int64, + device=next(model.parameters()).device, + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(active, op=dist.ReduceOp.SUM) + if int(active.item()) == 0: + continue + template_raw: object | None = raw + if dist.is_available() and dist.is_initialized(): + templates: list[object | None] = [None] * world_size + dist.all_gather_object(templates, raw) + template_raw = next( + (candidate for candidate in templates if candidate is not None), + None, + ) + if template_raw is None: + raise AlgorithmExecutionError( + f"Torch {split} evaluation has no typed batch template" + ) + template_batch = recipe_instance.adapt_batch(template_raw, batch_context) + if not isinstance(template_batch, TorchBatch): + raise AlgorithmConfigurationError( + "TorchRecipe evaluation batch template is invalid" + ) + while True: + local_active = raw is not None + active = torch.tensor( + 1 if local_active else 0, + dtype=torch.int64, + device=next(model.parameters()).device, + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(active, op=dist.ReduceOp.SUM) + if int(active.item()) == 0: + break + next_raw = next(iterator, None) if local_active else None + batch = ( + recipe_instance.adapt_batch(raw, batch_context) + if local_active + else zero_batch(template_batch) + ) + if not isinstance(batch, TorchBatch): + raise AlgorithmConfigurationError( + "TorchRecipe evaluation batch is invalid" + ) + with torch.no_grad(): + validation = recipe_instance.validation_step( + module_set, + batch, + TorchStepContext( + stage=stage_context, window_index=0, micro_batch_index=0 + ), + ) + if not isinstance(validation, TorchStepResult): + raise AlgorithmConfigurationError( + "TorchRecipe.validation_step returned invalid result" + ) + mapped_validation_metrics = map_metric_contributions(validation.metrics) + if ( + not isinstance(validation.loss, TorchCompositeLossContribution) + and set(mapped_validation_metrics) != declared_metric_names + ): + raise AlgorithmExecutionError( + "TorchRecipe.validation_step.metrics do not match TorchMetricPlan" + ) + for metric_name, contribution in aligned_metric_contributions( + mapped_validation_metrics + ).items(): + _accumulate_metric_totals( + evaluation_totals, + {metric_name: contribution}, + prefix=f"{split}_", + ) + if isinstance(validation.loss, TorchLossContribution): + totals = evaluation_totals.setdefault(f"{split}_loss", [0.0, 0.0]) + totals[0] += float(cast(Any, validation.loss.numerator).detach().item()) + totals[1] += float(validation.loss.normalizer) + elif isinstance(validation.loss, TorchCompositeLossContribution): + reduction = _reduce_composite_loss( + validation.loss, + config=config, + world_size=world_size, + device=next(model.parameters()).device, + dist=dist, + observation=reducer_observation, + expected_metrics=expected_metric_names, + ) + evaluation_metrics = { + name: contribution + for name, contribution in reduction.metrics.items() + if name != "train_loss" + } + _accumulate_metric_totals( + evaluation_totals, + evaluation_metrics, + prefix=f"{split}_", + ) + # The generic reducer's objective metric is reported under + # the established ``val_loss``/``test_loss`` name. + objective = reduction.metrics.get("train_loss") + if objective is not None: + _accumulate_metric_totals( + evaluation_totals, + {"loss": objective}, + prefix=f"{split}_", + ) + else: + raise AlgorithmExecutionError( + "TorchRecipe evaluation loss contribution is invalid" + ) + raw = next_raw + state_digest = hashlib.sha256( + json.dumps( + { + name: str(cast(Any, value).detach().cpu().numpy().tobytes()) + for name, value in model.state_dict().items() + }, + sort_keys=True, + ).encode() + ).hexdigest() + assigned = ray.get_runtime_context().get_assigned_resources() + resources = { + "num_cpus": float(assigned.get("CPU", assigned.get("cpu", 1.0))), + "num_gpus": float(assigned.get("GPU", assigned.get("gpu", 0.0))), + "custom": { + str(name): float(value) + for name, value in assigned.items() + if str(name).upper() not in {"CPU", "GPU"} + }, + } + node_id = ray.get_runtime_context().get_node_id() + input_rows_evidence = dict(coverage_totals) + input_rows_evidence[stage_roles[0]] = rows + execution_worker = { + "worker_id": f"torch-{rank}", + "node_id": str(node_id), + "rank": rank, + "world_size": world_size, + "shard_id": f"train-{rank}", + "resources": resources, + "model_state_digest": state_digest, + "rows_processed": rows, + "input_rows": input_rows_evidence, + "batch_count": steps, + "collective_steps": steps, + } + worker_records_raw: list[dict[str, object] | None] = [None] * world_size + if dist.is_available() and dist.is_initialized(): + dist.all_gather_object(worker_records_raw, execution_worker) + else: + worker_records_raw = [execution_worker] + if any(not isinstance(item, Mapping) for item in worker_records_raw): + raise AlgorithmExecutionError("Torch worker evidence collective is incomplete") + worker_records = [cast(dict[str, object], item) for item in worker_records_raw] + from ray.train import Checkpoint + + checkpoint_dir = Path(tempfile.mkdtemp(prefix="tributo_torch_checkpoint_")) + try: + torch.save( + getattr(model, "module", model).state_dict(), checkpoint_dir / "model.pt" + ) + torch.save(cast(Any, optimizer).state_dict(), checkpoint_dir / "optimizer.pt") + torch.save(cast(Any, scaler).state_dict(), checkpoint_dir / "scaler.pt") + if optimization.scheduler is not None: + torch.save( + cast(Any, optimization.scheduler).state_dict(), + checkpoint_dir / "scheduler.pt", + ) + final_cursor_values: list[object] = [None] * world_size + if dist.is_available() and dist.is_initialized(): + dist.all_gather_object(final_cursor_values, epoch_micro_batch_cursor) + else: + final_cursor_values = [epoch_micro_batch_cursor] + if any( + not isinstance(value, int) or isinstance(value, bool) or value < 0 + for value in final_cursor_values + ): + raise AlgorithmExecutionError("Torch final cursor collective is incomplete") + final_local_statistics = TorchRankProgressStatistics( + rows_processed=rows, + coverage_totals=coverage_totals, + loss_numerator_total=loss_numerator_total, + loss_normalizer_total=loss_normalizer_total, + metric_totals={ + name: (values[0], values[1]) for name, values in metric_totals.items() + }, + evaluation_totals={ + name: (values[0], values[1]) + for name, values in evaluation_totals.items() + }, + reducer_observation=reducer_observation, + ) + final_statistics_by_rank: list[object] = [None] * world_size + if dist.is_available() and dist.is_initialized(): + dist.all_gather_object( + final_statistics_by_rank, final_local_statistics.to_dict() + ) + else: + final_statistics_by_rank = [final_local_statistics.to_dict()] + if any(not isinstance(value, Mapping) for value in final_statistics_by_rank): + raise AlgorithmExecutionError( + "Torch final statistics collective is incomplete" + ) + final_progress = TorchCheckpointProgress( + epoch=max(epochs - 1, 0), + micro_batch_cursor=epoch_micro_batch_cursor, + optimizer_step=restored_step + steps // accumulation, + scheduler_step=scheduler_steps, + accumulation_steps=accumulation, + dataset_cursor_by_rank={ + str(rank_id): cast(int, cursor) + for rank_id, cursor in enumerate(final_cursor_values) + }, + shuffle_seed=int(seed + max(epochs - 1, 0)), + rows_processed=rows, + coverage_totals=coverage_totals, + loss_numerator_total=loss_numerator_total, + loss_normalizer_total=loss_normalizer_total, + metric_totals={ + name: (values[0], values[1]) for name, values in metric_totals.items() + }, + evaluation_totals={ + name: (values[0], values[1]) + for name, values in evaluation_totals.items() + }, + rank_statistics={ + str(rank_id): TorchRankProgressStatistics.from_dict( + cast(Mapping[str, Any], stats) + ) + for rank_id, stats in enumerate(final_statistics_by_rank) + }, + epoch_scheduler_applied=True, + ) + (checkpoint_dir / "torch_progress.json").write_text( + json.dumps(final_progress.to_dict(), sort_keys=True, separators=(",", ":")), + encoding="utf-8", + ) + rng_payload = torch.get_rng_state().cpu().numpy().tobytes() + rng_states: list[bytes | None] = [None] * world_size + if dist.is_available() and dist.is_initialized(): + dist.all_gather_object(rng_states, rng_payload) + else: + rng_states = [rng_payload] + if any(not isinstance(value, bytes) for value in rng_states): + raise AlgorithmExecutionError("Torch RNG state collective is incomplete") + cuda_rng_payload: list[list[bytes]] = [] + if torch.cuda.is_available(): + local_cuda = [ + state.cpu().numpy().tobytes() + for state in torch.cuda.get_rng_state_all() + ] + if dist.is_available() and dist.is_initialized(): + all_cuda: list[object] = [None] * world_size + dist.all_gather_object(all_cuda, local_cuda) + if any( + not isinstance(value, list) + or any(not isinstance(state, bytes) for state in value) + for value in all_cuda + ): + raise AlgorithmExecutionError( + "Torch CUDA RNG state collective is incomplete" + ) + cuda_rng_payload = cast(list[list[bytes]], all_cuda) + else: + cuda_rng_payload = [local_cuda] + torch.save( + { + "world_size": world_size, + "states": rng_states, + "cuda_states_by_rank": cuda_rng_payload, + }, + checkpoint_dir / "rng_state.pt", + ) + identity = stage_context.runtime.run_identity + if identity is None: + raise AlgorithmExecutionError( + "Torch Worker stage context has no run identity" + ) + payload_names = [ + "model.pt", + "optimizer.pt", + "scaler.pt", + "rng_state.pt", + "torch_progress.json", + ] + if optimization.scheduler is not None: + payload_names.append("scheduler.pt") + payload_files = { + name: hashlib.sha256((checkpoint_dir / name).read_bytes()).hexdigest() + for name in payload_names + } + descriptor = TorchCheckpointDescriptor( + schema_version=1, + identity=identity, + run_config_name=torch_run_config_name(identity), + state_layout=str(config.get("_core_state_layout", "replicated")), + world_size=world_size, + completed_step=restored_step + steps // accumulation, + policy_digest=stage_context.runtime.policy_digest, + execution_plan_digest=stage_context.runtime.execution_plan_digest, + input_binding_digest=str(config.get("_core_input_binding_digest", "")), + implementation_code_digest=str( + config.get("_core_implementation_code_digest", "") + ), + payload_files=payload_files, + adapter_identity=config.get("_core_adapter_identity"), + resume_supported=stage_context.runtime.resume_supported, + same_world_size_resume=stage_context.runtime.same_world_size_resume, + ) + + class _Draft: + checkpoint_dir: str | os.PathLike[str] + + def __init__( + self, checkpoint_dir: Path, checkpoint_owner_rank: int + ) -> None: + self.checkpoint_dir = str(checkpoint_dir) + self._checkpoint_owner_rank = checkpoint_owner_rank + + def report( + self, + *, + metrics: Mapping[str, object], + stage_context: object, + completed_step: int, + ) -> None: + del stage_context, completed_step + checkpoint = ( + Checkpoint.from_directory(str(checkpoint_dir)) + if rank == self._checkpoint_owner_rank + else None + ) + ray.train.report(dict(metrics), checkpoint=checkpoint) + + loss_state = torch.tensor( + [loss_numerator_total, loss_normalizer_total], + dtype=torch.float64, + device=next(model.parameters()).device, + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(loss_state, op=dist.ReduceOp.SUM) + train_loss = ( + float(loss_state[0].item() / loss_state[1].item()) + if loss_state[1].item() > 0 + else 0.0 + ) + metric_values: dict[str, float] = {} + metric_reducers = config.get( + "_core_metric_reducers", {"train_loss": "sum_count"} + ) + if not isinstance(metric_reducers, Mapping): + raise AlgorithmConfigurationError("Torch metric reducers must be a mapping") + metric_values.update( + _reduce_metric_totals( + metric_totals, + metric_reducers, + device=next(model.parameters()).device, + dist=dist, + world_size=world_size, + ) + ) + if composite_loss_seen and "train_loss" not in metric_values: + raise AlgorithmExecutionError( + "Composite Torch training produced no train_loss metric" + ) + metric_values.update( + _reduce_metric_totals( + evaluation_totals, + metric_reducers, + device=next(model.parameters()).device, + dist=dist, + world_size=world_size, + ) + ) + metric_values.setdefault("train_loss", train_loss) + for source_name, target_name in metric_mapping.items(): + if source_name in metric_values: + metric_values[target_name] = metric_values.pop(source_name) + reducer_report = { + "reducer_id": config.get("_core_global_loss_reducer_id"), + "reducer_api_version": config.get("_core_global_loss_reducer_api_version"), + "reducer_schema_id": config.get("_core_global_loss_reducer_schema_id"), + "reducer_code_digest": config.get("_core_global_loss_reducer_code_digest"), + "reducer_branch": reducer_observation.get("branch"), + "reducer_evidence": reducer_observation.get("evidence"), + } + reducer_report = { + key: value for key, value in reducer_report.items() if value is not None + } + report_torch_checkpoint( + { + "train_loss": train_loss, + **metric_values, + "execution_workers": worker_records, + "model_state_digest": state_digest, + "checkpoint_descriptor": descriptor.to_dict(), + **reducer_report, + }, + _Draft( + checkpoint_dir, + int(config.get("_core_checkpoint_owner_rank", 0)), + ), + stage_context, + restored_step + steps // accumulation, + ) + finally: + import shutil + + shutil.rmtree(checkpoint_dir, ignore_errors=True) + + +@DeveloperAPI +def torch_recipe_train_loop_per_worker(config: Mapping[str, Any]) -> None: + """Public Core worker entrypoint referenced by ``TorchStageSpec``.""" + _recipe_worker(config) + + +@DeveloperAPI +def ray_torch_adapter_train_loop_per_worker(config: Mapping[str, Any]) -> None: + """Public Core wrapper entrypoint for Adapter-owned Stage loops.""" + reference = config.get("_core_implementation_ref") + if not isinstance(reference, str): + raise AlgorithmConfigurationError( + "Adapter Worker implementation reference is missing" + ) + implementation = _load_reference(QualifiedReference.parse(reference)) + if not isinstance(implementation, type): + raise AlgorithmConfigurationError("Adapter Worker reference is invalid") + adapter = implementation() + if not isinstance(adapter, RayTorchAdapter): + raise AlgorithmConfigurationError("Adapter Worker reference is invalid") + stage_context_value = config.get("_core_stage_context") + if not isinstance(stage_context_value, Mapping): + raise AlgorithmConfigurationError("Adapter Worker stage context is missing") + stage_context = TorchStageContext.from_dict(stage_context_value) + checkpoint_context = _select_worker_checkpoint(config, stage_context) + try: + adapter.train_loop_per_worker( + config.get("adapter_config", {}), + checkpoint_context, + ) + finally: + if checkpoint_context.checkpoint is not None: + checkpoint_context.checkpoint.close() + + +@DeveloperAPI +class RayTrainTorchRuntime: + """Unified Core-owned Runtime for Recipe and Adapter implementations.""" + + @property + def runtime_id(self) -> str: + return RAY_TRAIN_TORCH_RUNTIME_ID + + def preflight( + self, + plan: Any, + run_id: str, + invocation_id: str, + ) -> TorchPreflightLease: + policy = _policy(plan) + if policy.state_layout == "sharded": + raise AlgorithmConfigurationError( + "Torch sharded state is reserved and not supported by Runtime v1" + ) + if policy.evidence_adapter_ref is not None: + raise AlgorithmConfigurationError( + "Torch evidence_adapter_ref is reserved until a Core evidence adapter protocol is gated" + ) + ray_config = plan.algorithm_config.get("ray", {}) + resume_config = ( + ray_config.get("resume", {}) if isinstance(ray_config, Mapping) else {} + ) + has_external_recovery = plan.runtime.resume_from is not None or ( + isinstance(resume_config, Mapping) + and any( + resume_config.get(name) is not None + for name in ("uri", "checkpoint_uri", "checkpoint_descriptor_digest") + ) + ) + if plan.runtime.torch_recovery is not None: + recovery = TorchRecoveryEnvelope.from_dict(plan.runtime.torch_recovery) + has_external_recovery = has_external_recovery or bool( + recovery.stage_checkpoints or recovery.active_checkpoint is not None + ) + if has_external_recovery and not policy.resume_supported: + raise AlgorithmConfigurationError( + "Torch Policy does not support external recovery" + ) + implementation = _load_torch_implementation(plan) + for stage in policy.execution_plan.stages: + expected_loop_ref = ( + _CORE_ADAPTER_LOOP_REF + if policy.loop_owner == "adapter" + else _CORE_RECIPE_LOOP_REF + ) + if stage.worker_loop_ref != expected_loop_ref: + raise AlgorithmConfigurationError( + "Torch Stage worker_loop_ref must use the Core-owned loop wrapper" + ) + stage_worker = _load_reference( + QualifiedReference.parse(stage.worker_loop_ref) + ) + if not callable(stage_worker): + raise AlgorithmConfigurationError( + f"Torch Stage {stage.stage_id!r} worker_loop_ref is not callable" + ) + identity = _identity( + plan, run_id, invocation_id, policy.execution_plan.final_stage_id + ) + context = TorchRuntimeContext( + algorithm_config=_torch_algorithm_context_config(plan), + implementation_id=plan.implementation.implementation_id, + world_size=plan.runtime.worker_count, + policy_digest=policy.digest, + execution_plan_digest=policy.execution_plan.digest, + run_identity=identity, + input_bindings=_torch_input_bindings(plan), + output_config=_torch_output_config(plan), + input_binding_digest=_input_binding_digest(plan), + state_layout=policy.state_layout, + adapter_identity=( + plan.implementation.implementation_id + if isinstance(implementation, RayTorchAdapter) + else None + ), + resume_supported=policy.resume_supported, + same_world_size_resume=policy.same_world_size_resume, + ) + if isinstance(implementation, RayTorchAdapter): + implementation.validate_environment(context) + metric_plan = implementation.metric_plan(context) + else: + self._validate_recipe_environment(context) + metric_plan = implementation.metric_plan(context) + if not isinstance(metric_plan, TorchMetricPlan): + raise AlgorithmConfigurationError( + "Torch metric_plan must return TorchMetricPlan" + ) + declared_reducers = { + name: reduction.value for name, reduction in policy.metric_reducers.items() + } + if dict(metric_plan.reducers) != declared_reducers: + raise AlgorithmConfigurationError( + "Torch metric_plan reducers do not match TorchPolicy" + ) + if policy.global_loss_reducer_ref: + reducer_reference = QualifiedReference.parse(policy.global_loss_reducer_ref) + _validate_module_digest( + reducer_reference, + policy.global_loss_reducer_code_digest, + ) + reducer = _load_reference(reducer_reference) + if isinstance(reducer, type): + reducer = reducer() + if not callable(getattr(reducer, "reduce", None)): + raise AlgorithmConfigurationError( + "Torch global loss reducer is not callable" + ) + if ( + getattr(reducer, "api_version", None) + != policy.global_loss_reducer_api_version + ): + raise AlgorithmConfigurationError( + "Torch global reducer API version mismatch" + ) + if ( + getattr(reducer, "component_schema_id", None) + != policy.composite_loss_schema_id + ): + raise AlgorithmConfigurationError( + "Torch global reducer schema mismatch" + ) + if ( + getattr(reducer, "code_digest", None) + != policy.global_loss_reducer_code_digest + ): + raise AlgorithmConfigurationError( + "Torch global reducer code digest mismatch" + ) + token = TorchPreflightTokenData( + run_id=run_id, + invocation_id=invocation_id, + algorithm=plan.resolution.algorithm, + implementation_ref=str(plan.implementation.implementation_ref), + implementation_code_digest=cast(str, plan.implementation.code_digest), + policy_digest=policy.digest, + execution_plan_digest=policy.execution_plan.digest, + runtime_id=self.runtime_id, + reducer_identity=policy.global_loss_reducer_ref, + plan_digest=plan.plan_id, + ) + return TorchPreflightLease(token) + + @staticmethod + def _validate_recipe_environment(context: TorchRuntimeContext) -> None: + if context.world_size < 1: + raise AlgorithmConfigurationError("Torch world size must be positive") + try: + import ray + import ray.train + import torch + except ImportError as exc: + raise AlgorithmConfigurationError( + "TorchRecipe requires Ray and PyTorch" + ) from exc + if not hasattr(ray, "train") or not hasattr(torch, "nn"): + raise AlgorithmConfigurationError("TorchRecipe environment is incomplete") + + def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionResult: + if not isinstance(envelope, TorchRuntimeExecutionEnvelope): + raise AlgorithmConfigurationError( + "Ray Train Torch requires TorchRuntimeExecutionEnvelope" + ) + base = envelope.base + if base.cancelled: + raise AlgorithmExecutionError("Torch execution was cancelled") + invocation_id = envelope.preflight_lease.data.invocation_id + token = envelope.preflight_lease.consume( + run_id=base.run_id, + invocation_id=invocation_id, + plan_digest=base.plan.plan_id, + runtime_id=self.runtime_id, + ) + implementation = _load_torch_implementation(base.plan) + policy = _policy(base.plan) + if policy.state_layout == "sharded": + raise AlgorithmConfigurationError( + "Torch sharded state is reserved and not supported by Runtime v1" + ) + if policy.evidence_adapter_ref is not None: + raise AlgorithmConfigurationError( + "Torch evidence_adapter_ref is reserved until a Core evidence adapter protocol is gated" + ) + reducer_metadata = _reducer_metadata(policy) + completed_stage_ids, active_stage_id, recovery_records = _recovery_records( + base.plan, + policy, + worker_count=base.plan.runtime.worker_count, + ) + prepared = _prepare_datasets(base) + try: + stages = policy.execution_plan.stages + stage_records: dict[str, dict[str, Any]] = dict(recovery_records) + last_result: Any = None + final_result: Any = None + stage_evidence: list[ComponentStageEvidence] = [] + final_expected_rows: Mapping[str, int] = {} + for index, stage in enumerate(stages): + if stage.stage_id in completed_stage_ids: + recovered_record = stage_records.get(stage.stage_id) + if recovered_record is None: + raise AlgorithmExecutionError( + f"Torch recovery is missing Stage {stage.stage_id!r}" + ) + recovered_evidence = _recovered_stage_evidence( + plan=base.plan, + policy=policy, + stage=stage, + descriptor=recovered_record["descriptor"], + evidence=recovered_record.get("evidence", {}), + ) + if policy.state_layout == "component": + stage_evidence.append(recovered_evidence) + if stage.stage_id == policy.execution_plan.final_stage_id: + final_expected_rows = { + role.role: int(role.expected_rows or 0) + for role in recovered_evidence.roles + if role.present and role.expected_rows is not None + } + continue + identity = _identity( + base.plan, token.run_id, token.invocation_id, stage.stage_id + ) + runtime_context = TorchRuntimeContext( + algorithm_config=_torch_algorithm_context_config(base.plan), + implementation_id=base.plan.implementation.implementation_id, + world_size=base.plan.runtime.worker_count, + policy_digest=policy.digest, + execution_plan_digest=policy.execution_plan.digest, + run_identity=identity, + input_bindings=_torch_input_bindings(base.plan), + output_config=_torch_output_config(base.plan), + input_binding_digest=_input_binding_digest(base.plan), + state_layout=policy.state_layout, + adapter_identity=( + base.plan.implementation.implementation_id + if isinstance(implementation, RayTorchAdapter) + else None + ), + resume_supported=policy.resume_supported, + same_world_size_resume=policy.same_world_size_resume, + ) + predecessor_id = stage.checkpoint_from_stage + predecessor_record = ( + stage_records.get(predecessor_id) + if predecessor_id is not None + else None + ) + context_descriptor = ( + { + key: value + for key, value in predecessor_record["descriptor"].items() + if key not in {"locator", "checkpoint_locator"} + } + if predecessor_record is not None + else None + ) + context = _stage_context( + base.plan, + runtime_context, + stage, + index, + predecessor=predecessor_id, + predecessor_descriptor=context_descriptor, + ) + stage_datasets: Mapping[str, object] = prepared.views + if isinstance(implementation, RayTorchAdapter): + stage_datasets = implementation.bind_datasets( + prepared.views, + context, + ) + if not isinstance(stage_datasets, Mapping) or not stage_datasets: + raise AlgorithmConfigurationError( + "RayTorchAdapter.bind_datasets must return named datasets" + ) + expected_rows = _validate_stage_routes( + policy, + stage, + stage_datasets, + base.plan.runtime.worker_count, + ) + if stage.stage_id == policy.execution_plan.final_stage_id: + final_expected_rows = dict(expected_rows) + if ( + stage.checkpoint_from_stage is not None + and predecessor_record is None + ): + raise AlgorithmExecutionError( + f"Torch Stage {stage.stage_id!r} is missing checkpoint from " + f"{stage.checkpoint_from_stage!r}" + ) + core_control = _control_for_stage( + base.plan, + policy, + stage, + run_id=token.run_id, + invocation_id=token.invocation_id, + checkpoint=( + stage_records.get(stage.stage_id) + if active_stage_id == stage.stage_id + else stage_records.get(stage.checkpoint_from_stage) + if stage.checkpoint_from_stage is not None + else None + ), + purpose=( + "cross_run_initial_recovery" + if active_stage_id == stage.stage_id + else "stage_dependency" + if stage.checkpoint_from_stage is not None + else None + ), + source_stage_id=( + None + if active_stage_id == stage.stage_id + else stage.checkpoint_from_stage + ), + ) + train_config = _torch_algorithm_context_config(base.plan) + stage_binding = base.plan.input_bindings.get(stage.input_roles[0]) + if stage_binding is None: + raise AlgorithmConfigurationError( + f"Torch Stage role {stage.input_roles[0]!r} has no input binding" + ) + train_config.update( + { + "_core_implementation_ref": str( + base.plan.implementation.implementation_ref + ), + "_core_implementation_code_digest": base.plan.implementation.code_digest, + "_core_state_layout": policy.state_layout, + "_core_checkpoint_owner_rank": policy.checkpoint_owner_rank, + "_core_input_binding_digest": _input_binding_digest(base.plan), + "_core_feature_names": list(stage_binding.feature_names), + "_core_label_name": stage_binding.label_name, + "_core_weight_name": stage_binding.sample_weight_name, + "_core_stage_input_roles": list(stage.input_roles), + "_core_input_role_bindings": _torch_input_bindings(base.plan), + "_core_checkpoint_interval_windows": int( + getattr(stage, "checkpoint_interval_windows", 1) + ), + "_core_policy_digest": policy.digest, + "_core_execution_plan_digest": policy.execution_plan.digest, + "_core_global_loss_reducer_ref": policy.global_loss_reducer_ref, + "_core_global_loss_reducer_api_version": policy.global_loss_reducer_api_version, + "_core_global_loss_reducer_code_digest": policy.global_loss_reducer_code_digest, + "_core_composite_loss_schema_id": policy.composite_loss_schema_id, + "_core_global_loss_reducer_id": reducer_metadata.get( + "reducer_id" + ), + "_core_global_loss_reducer_schema_id": reducer_metadata.get( + "reducer_schema_id" + ), + "_core_metric_reducers": { + name: reduction.value + for name, reduction in policy.metric_reducers.items() + }, + "_core_metric_mapping": dict( + getattr(stage, "metric_mapping", {}) + ), + "_core_adapter_identity": ( + base.plan.implementation.implementation_id + if isinstance(implementation, RayTorchAdapter) + else None + ), + "_core_stage_context": context.to_dict(), + "_core_batch_size": int( + base.plan.algorithm_config.get("training", {}).get( + "batch_size", 32 + ) + if isinstance( + base.plan.algorithm_config.get("training", {}), Mapping + ) + else 32 + ), + "_core_torch_evidence": {}, + "core_control": core_control, + "_core_checkpoint_opener_ref": ( + "tributo.integrations.algorithm_runtimes.ray_train_torch:" + "open_torch_checkpoint_locator" + if core_control is not None + else None + ), + } + ) + loop: Any + if isinstance(implementation, TorchRecipe): + loop_ref = _load_reference( + QualifiedReference.parse(stage.worker_loop_ref) + ) + if not callable(loop_ref): + raise AlgorithmConfigurationError( + "Recipe Stage worker_loop_ref is not callable" + ) + loop = loop_ref + else: + adapter = cast(RayTorchAdapter, implementation) + adapter_config = adapter.worker_config(context) + if not isinstance(adapter_config, Mapping): + raise AlgorithmConfigurationError( + "RayTorchAdapter.worker_config must return a mapping" + ) + if any(not isinstance(key, str) for key in adapter_config): + raise AlgorithmConfigurationError( + "Adapter worker config keys must be strings" + ) + _validate_adapter_worker_config(adapter_config) + try: + json.dumps(adapter_config, allow_nan=False) + except (TypeError, ValueError) as exc: + raise AlgorithmConfigurationError( + "Adapter worker config must be JSON-compatible" + ) from exc + train_config["adapter_config"] = dict(adapter_config) + loop_ref = _load_reference( + QualifiedReference.parse(stage.worker_loop_ref) + ) + if not callable(loop_ref): + raise AlgorithmConfigurationError( + "Adapter Stage worker_loop_ref is not callable" + ) + loop = loop_ref + from ray.train import FailureConfig, RunConfig, ScalingConfig + from ray.train.torch import TorchConfig, TorchTrainer + + from tributo.integrations.algorithm_runtimes.ray_data_config import ( + TorchRoleDataConfig, + ) + + storage_path = None + ray_config = base.plan.algorithm_config.get("ray", {}) + max_failures = 0 + if isinstance(ray_config, Mapping): + storage_path = ray_config.get("storage_path") + configured_failures = ray_config.get("max_failures", 0) + if ( + not isinstance(configured_failures, int) + or isinstance(configured_failures, bool) + or configured_failures < -1 + ): + raise AlgorithmConfigurationError( + "ray.max_failures must be -1 or a non-negative integer" + ) + max_failures = configured_failures + if ( + storage_path is None + and cast(Any, base.plan.runtime.execution_profile).value == "local" + ): + storage_path = tempfile.mkdtemp(prefix="tributo_torch_runs_") + if storage_path is not None: + claim_torch_run_directory(storage_path, identity) + trainer = TorchTrainer( + train_loop_per_worker=loop, + train_loop_config=train_config, + scaling_config=ScalingConfig( + num_workers=base.plan.runtime.worker_count, + use_gpu=base.plan.runtime.num_gpus > 0, + resources_per_worker=_resource_map(base.plan), + placement_strategy="SPREAD", + ), + datasets=cast(Any, dict(stage_datasets)), + run_config=RunConfig( + name=torch_run_config_name(identity), + storage_path=str(storage_path) + if storage_path is not None + else None, + failure_config=FailureConfig(max_failures=max_failures), + ), + torch_config=TorchConfig( + backend=None if policy.backend == "auto" else policy.backend + ), + dataset_config=TorchRoleDataConfig( + {route.role: route for route in policy.dataset_routing} + ), + ) + last_result = trainer.fit() + if stage.stage_id == policy.execution_plan.final_stage_id: + final_result = last_result + metrics = last_result.metrics or {} + stage_descriptor_payload = metrics.get("checkpoint_descriptor") + checkpoint = getattr(last_result, "checkpoint", None) + if checkpoint is not None and isinstance( + stage_descriptor_payload, Mapping + ): + validated_descriptor = describe_torch_checkpoint( + TorchCheckpointRef(checkpoint), + TorchCheckpointContext( + stage=context, + run_id=identity.run_id, + invocation_id=identity.invocation_id, + checkpoint_owner="core", + ), + ) + if validated_descriptor.to_dict() != dict(stage_descriptor_payload): + raise AlgorithmExecutionError( + "Torch Stage checkpoint descriptor differs from embedded payload" + ) + persisted_locator = _persist_stage_checkpoint( + checkpoint, + identity=identity, + storage_path=storage_path, + descriptor_digest=validated_descriptor.digest, + ) + if persisted_locator is not None: + stage_descriptor_payload = dict(stage_descriptor_payload) + stage_descriptor_payload["locator"] = persisted_locator + stage_descriptor_payload["descriptor_digest"] = ( + validated_descriptor.digest + ) + stage_records[stage.stage_id] = { + "locator": ( + stage_descriptor_payload.get("locator") + if isinstance(stage_descriptor_payload, Mapping) + else None + ), + "descriptor_digest": validated_descriptor.digest, + "descriptor": validated_descriptor.to_dict(), + "evidence": { + key: metrics[key] + for key in ( + "execution_workers", + "model_state_digest", + "reducer_id", + "reducer_api_version", + "reducer_schema_id", + "reducer_code_digest", + "reducer_branch", + "reducer_evidence", + ) + if key in metrics + }, + } + elif stage.checkpoint_required: + raise AlgorithmExecutionError( + f"Torch Stage {stage.stage_id!r} requires a Core checkpoint" + ) + stage_evidence.append( + _component_stage_evidence( + plan=base.plan, + policy=policy, + stage=stage, + identity=identity, + metrics=metrics, + expected_rows=expected_rows, + ) + ) + final_stage = next( + stage + for stage in stages + if stage.stage_id == policy.execution_plan.final_stage_id + ) + if final_result is None and final_stage.stage_id in stage_records: + recovered = stage_records[final_stage.stage_id] + locator = TorchCheckpointLocator( + cast(str, recovered["locator"]), + cast(str, recovered["descriptor_digest"]), + ) + recovered_checkpoint = open_torch_checkpoint_locator(locator) + final_result = _CheckpointResultProxy( + recovered_checkpoint, + recovered_checkpoint, + metrics={ + **dict(recovered.get("evidence", {})), + "checkpoint_descriptor": recovered["descriptor"], + }, + ) + if final_result is None: + raise AlgorithmExecutionError( + "Torch execution plan has no Stage result" + ) + last_result = final_result + metrics = dict(final_result.metrics or {}) + checkpoint = getattr(final_result, "checkpoint", None) + if checkpoint is not None: + raw_descriptor = metrics.get("checkpoint_descriptor") + if not isinstance(raw_descriptor, Mapping): + raise AlgorithmExecutionError( + "Torch final Checkpoint is missing its Core descriptor" + ) + descriptor = TorchCheckpointDescriptor.from_dict(raw_descriptor) + checkpoint_ref = TorchCheckpointRef( + checkpoint=checkpoint, + descriptor_digest=descriptor.digest, + source_stage_id=descriptor.identity.stage_id, + descriptor=descriptor, + ) + describe_torch_checkpoint( + checkpoint_ref, + TorchCheckpointContext( + stage=_stage_context( + base.plan, + TorchRuntimeContext( + algorithm_config=_torch_algorithm_context_config( + base.plan + ), + implementation_id=base.plan.implementation.implementation_id, + world_size=base.plan.runtime.worker_count, + policy_digest=policy.digest, + execution_plan_digest=policy.execution_plan.digest, + run_identity=descriptor.identity, + input_bindings=_torch_input_bindings(base.plan), + output_config=_torch_output_config(base.plan), + input_binding_digest=_input_binding_digest(base.plan), + state_layout=policy.state_layout, + adapter_identity=( + base.plan.implementation.implementation_id + if isinstance(implementation, RayTorchAdapter) + else None + ), + resume_supported=policy.resume_supported, + same_world_size_resume=policy.same_world_size_resume, + ), + final_stage, + stages.index(final_stage), + ), + run_id=descriptor.identity.run_id, + invocation_id=descriptor.identity.invocation_id, + checkpoint_owner="core", + ), + ) + final_identity = _identity( + base.plan, + token.run_id, + token.invocation_id, + policy.execution_plan.final_stage_id, + ) + supplied_evidence = metrics.get("torch_evidence") + if supplied_evidence is not None: + raise AlgorithmExecutionError( + "Torch execution evidence is Core-owned and must not be supplied by an algorithm" + ) + raw_workers = metrics.get("execution_workers") + if ( + not isinstance(raw_workers, (list, tuple)) + or len(raw_workers) != base.plan.runtime.worker_count + ): + raise AlgorithmExecutionError( + "Torch execution did not report every worker" + ) + workers = tuple( + WorkerExecutionEvidence.from_dict(item) + for item in _normalize_worker_evidence(raw_workers, base.plan) + ) + role_evidence = _role_execution_evidence( + plan=base.plan, + policy=policy, + stage=final_stage, + workers=workers, + expected_rows=final_expected_rows, + ) + global_digest = metrics.get("model_state_digest") + if not isinstance(global_digest, str) or len(global_digest) != 64: + raise AlgorithmExecutionError( + "Torch execution did not report a model digest" + ) + composition_digest = None + if policy.state_layout == "component": + composition_digest = hashlib.sha256( + json.dumps( + [item.to_dict() for item in stage_evidence], + sort_keys=True, + separators=(",", ":"), + ).encode() + ).hexdigest() + metrics["torch_evidence"] = TorchExecutionEvidence( + identity=final_identity, + run_config_name=torch_run_config_name(final_identity), + policy_digest=policy.digest, + parallelism_id=policy.parallelism_id, + state_layout=policy.state_layout, + workers=workers, + roles=role_evidence, + replicated_state=( + ReplicatedTorchStateEvidence( + model_digests_by_rank={ + worker.rank: cast(str, worker.model_state_digest) + for worker in workers + }, + global_model_digest=global_digest, + ) + if policy.state_layout == "replicated" + else None + ), + stages=tuple(stage_evidence) + if policy.state_layout == "component" + else (), + composition_digest=composition_digest, + final_stage_id=( + policy.execution_plan.final_stage_id + if policy.state_layout == "component" + else None + ), + reducer_id=cast(str | None, reducer_metadata.get("reducer_id")), + reducer_api_version=cast( + int | None, reducer_metadata.get("reducer_api_version") + ), + reducer_schema_id=cast( + str | None, reducer_metadata.get("reducer_schema_id") + ), + reducer_code_digest=cast( + str | None, reducer_metadata.get("reducer_code_digest") + ), + reducer_branch=( + metrics.get("reducer_branch") + if isinstance(metrics.get("reducer_branch"), str) + else None + ), + reducer_evidence=( + dict(metrics["reducer_evidence"]) + if isinstance(metrics.get("reducer_evidence"), Mapping) + else {} + ), + ).to_dict() + execution = AlgorithmExecutionResult( + status="succeeded", + metrics=portable_fit_only_metrics( + metrics, extra_evidence_names=_TORCH_INTERNAL_METRIC_NAMES + ), + ) + export_result = _CheckpointResultProxy( + last_result, + checkpoint, + metrics=metrics, + core_evidence_attested=True, + ) + export_state_details: dict[str, object] = {} + if policy_result_policy(base.plan) is ResultPolicy.BUNDLE_REQUIRED: + if isinstance(implementation, RayTorchAdapter): + final_context = TorchCheckpointContext( + stage=_stage_context( + base.plan, + TorchRuntimeContext( + algorithm_config=_torch_algorithm_context_config( + base.plan + ), + implementation_id=base.plan.implementation.implementation_id, + world_size=base.plan.runtime.worker_count, + policy_digest=policy.digest, + execution_plan_digest=policy.execution_plan.digest, + run_identity=final_identity, + input_bindings=_torch_input_bindings(base.plan), + output_config=_torch_output_config(base.plan), + input_binding_digest=_input_binding_digest(base.plan), + state_layout=policy.state_layout, + adapter_identity=( + base.plan.implementation.implementation_id + if isinstance(implementation, RayTorchAdapter) + else None + ), + resume_supported=policy.resume_supported, + same_world_size_resume=policy.same_world_size_resume, + ), + final_stage, + stages.index(final_stage), + ), + run_id=token.run_id, + invocation_id=token.invocation_id, + checkpoint_owner="core", + ) + checkpoint = implementation.checkpoint_source( + last_result, final_context + ) + if checkpoint is None: + raise AlgorithmExecutionError( + "RayTorchAdapter checkpoint_source returned no checkpoint" + ) + checkpoint_ref = ( + checkpoint + if isinstance(checkpoint, TorchCheckpointRef) + else TorchCheckpointRef(checkpoint) + ) + describe_torch_checkpoint(checkpoint_ref, final_context) + export_result = _CheckpointResultProxy( + last_result, + checkpoint_ref.checkpoint, + metrics=metrics, + core_evidence_attested=True, + ) + execution = export_ray_train_torch_result( + result=export_result, + plan=base.plan, + run_id=token.run_id, + state_details_sink=export_state_details, + ) + raw_worker_metadata = metrics.get("execution_workers", []) + normalized_worker_metadata = ( + [ + dict(item) + for item in _normalize_worker_evidence( + raw_worker_metadata, base.plan + ) + ] + if isinstance(raw_worker_metadata, (list, tuple)) + else [] + ) + state_details = ( + _component_state_details(tuple(stage_evidence)) + if policy.state_layout == "component" + else export_state_details + ) + return WorkerExecutionResult( + execution=execution, + actual_versions=_actual_environment_versions( + base.plan.environment.python, + base.plan.environment.dependencies, + ), + worker_metadata={ + "topology": "ray_train_torch", + "workers": normalized_worker_metadata, + "state": { + "coordination": "torch_managed", + "synchronized": True, + "bounded": True, + "global_model_digest": metrics.get("model_state_digest"), + "details": state_details, + }, + "torch_evidence": metrics.get("torch_evidence", {}), + "input_complete": True, + "driver_materialized_training_rows": 0, + }, + ) + finally: + prepared.close() + + +@DeveloperAPI +def create_torch_algorithm( + *, plan: Any, implementation: object, artifacts: tuple[object, ...] +) -> object: + """Factory validation hook retained for the unified Builder.""" + del artifacts + expected = _load_torch_implementation(plan) + if implementation is not _load_reference(plan.implementation.implementation_ref): + raise AlgorithmConfigurationError( + "Torch implementation drifted after descriptor resolution" + ) + return expected + + +class _CheckpointResultProxy: + """Preserve Ray Result metrics while replacing the exported Checkpoint.""" + + def __init__( + self, + result: object, + checkpoint: object, + *, + metrics: Mapping[str, object] | None = None, + core_evidence_attested: bool = False, + ) -> None: + self.metrics = dict(metrics or getattr(result, "metrics", {}) or {}) + self.checkpoint = checkpoint + self.core_evidence_attested = core_evidence_attested + + +def _component_composition_digest(result: object) -> str | None: + """Read a validated component composition digest from a Torch result.""" + if getattr(result, "core_evidence_attested", False) is not True: + return None + metrics = getattr(result, "metrics", {}) or {} + evidence = metrics.get("torch_evidence") if isinstance(metrics, Mapping) else None + digest = ( + evidence.get("composition_digest") if isinstance(evidence, Mapping) else None + ) + if ( + not isinstance(digest, str) + or len(digest) != 64 + or any(char not in "0123456789abcdef" for char in digest) + ): + return None + return digest + + +def _source_state_details(metadata: Mapping[str, object]) -> dict[str, object]: + """Project Adapter-declared scalar export evidence into receipt state details.""" + details: dict[str, object] = {} + for source_name, detail_name in ( + ("sampling", "sampling"), + ("topology_kind", "topology_kind"), + ("sparse_routing", "routing"), + ): + value = metadata.get(source_name) + if isinstance(value, (str, int, float, bool)): + details[detail_name] = value + if "routing" in details: + details["jagged"] = True + return details + + +@DeveloperAPI +def export_ray_train_torch_result( + *, + result: object, + plan: Any, + run_id: str, + state_details_sink: dict[str, object] | None = None, +) -> AlgorithmExecutionResult: + """Export the final Stage through the Core BundleExportService.""" + if ( + plan.distribution_spec is None + or plan.distribution_spec.result_policy is ResultPolicy.FIT_ONLY + ): + return AlgorithmExecutionResult( + status="succeeded", + metrics=portable_fit_only_metrics( + getattr(result, "metrics", {}) or {}, + extra_evidence_names=_TORCH_INTERNAL_METRIC_NAMES, + ), + ) + import importlib.metadata + + from tributo.exporting.models import BundleOutputConfig, ExportTarget + from tributo.exporting.service import BundleExportService + from tributo.integrations.sources.ray_torch import ( + RayTorchSourceProvider, + TorchSourceOptions, + ) + + output = plan.algorithm_config.get("output", {}) + if not isinstance(output, Mapping) or not isinstance(output.get("bundle_uri"), str): + raise AlgorithmConfigurationError( + "Torch Bundle export requires output.bundle_uri" + ) + composition_digest = None + if _policy(plan).state_layout == "component": + composition_digest = _component_composition_digest(result) + if composition_digest is None: + raise AlgorithmExecutionError( + "Torch component Bundle is missing composition_digest" + ) + final_stage_id = getattr(_policy(plan).execution_plan, "final_stage_id", None) + final_stage = next( + ( + stage + for stage in _policy(plan).execution_plan.stages + if getattr(stage, "stage_id", final_stage_id) == final_stage_id + ), + _policy(plan).execution_plan.stages[-1], + ) + options = TorchSourceOptions( + implementation_ref=str(plan.implementation.implementation_ref), + implementation_code_digest=plan.implementation.code_digest, + implementation_id=plan.implementation.implementation_id, + loop_owner=_policy(plan).loop_owner, + algorithm_config=_torch_algorithm_context_config(plan), + input_bindings=_torch_input_bindings(plan), + output_config=_torch_output_config(plan), + policy_digest=_policy(plan).digest, + plan_digest=plan.plan_id, + input_binding_digest=_input_binding_digest(plan), + stage_input_roles=tuple(final_stage.input_roles), + stage_index=_policy(plan).execution_plan.stages.index(final_stage), + ) + provider = RayTorchSourceProvider() + with provider.open_source(result, options) as source: + if state_details_sink is not None: + state_details_sink.update(_source_state_details(source.metadata)) + artifact_payload = source.metadata.get("artifact_plan", {}) + if not isinstance(artifact_payload, Mapping): + raise AlgorithmExecutionError("Torch artifact plan is missing") + raw_targets = artifact_payload.get("targets") + raw_roles = artifact_payload.get("roles") + if not isinstance(raw_targets, (list, tuple)) or not raw_targets: + raise AlgorithmExecutionError("Torch artifact plan declares no targets") + if not isinstance(raw_roles, Mapping) or not raw_roles: + raise AlgorithmExecutionError( + "Torch artifact plan declares no Bundle roles" + ) + targets = tuple( + ExportTarget(**dict(target)) + for target in raw_targets + if isinstance(target, Mapping) + ) + if len(targets) != len(raw_targets): + raise AlgorithmExecutionError("Torch artifact plan target is malformed") + published = BundleExportService().export_bundle( + source, + BundleOutputConfig( + bundle_uri=output["bundle_uri"], + request_id=run_id, + run_id=run_id, + targets=list(targets), + roles=dict(raw_roles), + ), + tributo_version=importlib.metadata.version("tributo"), + ) + outputs: dict[str, object] = { + "bundle_id": published.bundle_id, + "bundle_uri": published.canonical_uri, + "execution_id": published.execution_id, + "manifest_sha256": published.manifest_sha256, + } + if composition_digest is not None: + outputs["composition_digest"] = composition_digest + return AlgorithmExecutionResult( + status="succeeded", + metrics=portable_fit_only_metrics( + getattr(result, "metrics", {}) or {}, + extra_evidence_names=_TORCH_INTERNAL_METRIC_NAMES, + ), + outputs=outputs, + ) + + +__all__ = [ + "RAY_TRAIN_TORCH_RUNTIME_ID", + "RayTrainTorchRuntime", + "ray_torch_adapter_train_loop_per_worker", + "torch_recipe_train_loop_per_worker", + "create_torch_algorithm", + "export_ray_train_torch_result", +] diff --git a/src/tributo/integrations/algorithm_runtimes/torch_recipe.py b/src/tributo/integrations/algorithm_runtimes/torch_recipe.py deleted file mode 100644 index 74cdb7a..0000000 --- a/src/tributo/integrations/algorithm_runtimes/torch_recipe.py +++ /dev/null @@ -1,1483 +0,0 @@ -"""Lower narrow PyTorch recipes onto Ray Train's collective runtime.""" - -from __future__ import annotations - -import json -import math -import random -import shutil -import tempfile -from collections.abc import Mapping -from pathlib import Path -from typing import Any, Protocol, runtime_checkable - -from pydantic import Field, model_validator - -from tributo._common.config import StrictConfigModel -from tributo.algorithms.api import ( - AlgorithmConfigurationError, - AlgorithmExecutionResult, - MetricReduction, - ResolvedAlgorithmPlan, -) -from tributo.algorithms.core.worker import _load_reference, _validate_module_digest -from tributo.algorithms.spi import ( - CollectiveAlgorithm, - MetricPlan, - OptimizationPlan, - TorchTrainingRecipe, - TrainingRecipeV2, - TrainingStepResult, -) -from tributo.util.annotations import DeveloperAPI - -_TRAINER_TYPE = "torch_recipe" - - -@runtime_checkable -class _FirstPartyTorchRecipeAdapter(Protocol): - """Internal typed contract for first-party domain-specific Recipe hooks.""" - - _trainer_type: str - - def _bind_datasets( - self, - datasets: Mapping[str, object], - *, - config: Mapping[str, Any], - worker_count: int, - resume_from: str | None, - ) -> Mapping[str, object]: ... - - def _lower_worker_config( - self, - config: Mapping[str, Any], - ) -> Mapping[str, Any]: ... - - def _prepare_batch( - self, - batch: object, - *, - feature_names: tuple[str, ...], - label_name: str, - weight_name: str | None, - config: Mapping[str, Any], - ) -> tuple[object, object, object | None, int]: ... - - def _checkpoint_contract( - self, - *, - config: Mapping[str, Any], - feature_count: int, - output_shape: tuple[int, ...], - framework_version: str, - model_digest: str, - world_size: int, - ) -> dict[str, Any]: ... - - def _write_checkpoint_artifacts(self, checkpoint_dir: Path) -> tuple[str, ...]: ... - - def _validate_checkpoint_artifacts(self, checkpoint_dir: Path) -> None: ... - - -class _LoopConfig(StrictConfigModel): - epochs: int = Field(default=1, ge=1) - batch_size: int = Field(default=256, ge=1) - prefetch_batches: int = Field(default=1, ge=0) - local_shuffle_buffer_size: int | None = Field(default=None, ge=1) - seed: int = 42 - amp: bool = False - early_stopping_patience: int | None = Field(default=None, ge=1) - - @model_validator(mode="after") - def validate_shuffle_buffer(self) -> _LoopConfig: - """Keep Ray's local shuffle buffer at least one full batch.""" - if ( - self.local_shuffle_buffer_size is not None - and self.local_shuffle_buffer_size < self.batch_size - ): - raise ValueError( - "local_shuffle_buffer_size must be at least training.batch_size" - ) - return self - - -class _OutputConfig(StrictConfigModel): - bundle_uri: str = Field(min_length=1) - - -def _mapping(value: object, name: str) -> Mapping[str, Any]: - if not isinstance(value, Mapping): - raise AlgorithmConfigurationError(f"{name} must be a mapping") - if any(not isinstance(key, str) for key in value): - raise AlgorithmConfigurationError(f"{name} keys must be strings") - return value - - -def _recipe_type( - reference: object, code_digest: str | None -) -> type[TorchTrainingRecipe] | type[TrainingRecipeV2]: - from tributo.algorithms.api import QualifiedReference - - if not isinstance(reference, QualifiedReference): - raise AlgorithmConfigurationError("recipe reference is invalid") - _validate_module_digest(reference, code_digest) - implementation = _load_reference(reference) - if not isinstance(implementation, type) or not issubclass( - implementation, (TorchTrainingRecipe, TrainingRecipeV2) - ): - raise AlgorithmConfigurationError( - "Torch recipe implementation must subclass TorchTrainingRecipe or " - "TrainingRecipeV2" - ) - expected_version = 2 if issubclass(implementation, TrainingRecipeV2) else 1 - if getattr(implementation, "api_version", None) != expected_version: - raise AlgorithmConfigurationError( - f"Torch recipe api_version must be {expected_version}" - ) - return implementation - - -def _new_recipe( - reference: object, - code_digest: str | None, -) -> TorchTrainingRecipe | TrainingRecipeV2: - recipe_cls = _recipe_type(reference, code_digest) - try: - return recipe_cls() - except TypeError as exc: - raise AlgorithmConfigurationError( - "Torch recipe classes must have a no-argument constructor" - ) from exc - - -class _TorchRecipeCollectiveAlgorithm(CollectiveAlgorithm): - """Internal adapter that owns infrastructure around one user recipe.""" - - def __init__( - self, - plan: ResolvedAlgorithmPlan, - recipe: TorchTrainingRecipe | TrainingRecipeV2, - ) -> None: - self._plan = plan - self._recipe = recipe - ray_config = _mapping(plan.algorithm_config.get("ray", {}), "ray config") - max_failures = ray_config.get("max_failures", 0) - if ( - not isinstance(max_failures, int) - or isinstance(max_failures, bool) - or max_failures != 0 - ): - raise AlgorithmConfigurationError( - "Torch recipes require ray.max_failures=0 until the independent " - "failure-recovery gate passes" - ) - resume = _mapping(ray_config.get("resume", {}), "ray.resume config") - configured_checkpoint = resume.get("checkpoint_path") - if configured_checkpoint is not None and ( - configured_checkpoint != plan.runtime.resume_from - ): - raise AlgorithmConfigurationError( - "algorithm_config.ray.resume.checkpoint_path must match " - "ExecutionRequest.resume_from" - ) - - def bind_datasets(self, datasets: Mapping[str, object]) -> Mapping[str, object]: - """Accept pre-bound train/val/test datasets without performing a split.""" - if isinstance(self._recipe, _FirstPartyTorchRecipeAdapter): - bound = self._recipe._bind_datasets( - datasets, - config=self._plan.algorithm_config, - worker_count=self._plan.runtime.worker_count, - resume_from=self._plan.runtime.resume_from, - ) - if not isinstance(bound, Mapping) or not bound: - raise AlgorithmConfigurationError( - "Torch recipe dataset adapter must return named Datasets" - ) - return dict(bound) - if "train" in datasets: - unknown = sorted(set(datasets) - {"train", "val", "test"}) - if unknown: - raise AlgorithmConfigurationError( - f"Torch recipe received unknown dataset role(s): {unknown}" - ) - return dict(datasets) - if len(datasets) != 1: - raise AlgorithmConfigurationError( - "Torch recipe input must bind one train Dataset or explicit " - "train/val/test Dataset roles" - ) - return {"train": next(iter(datasets.values()))} - - def build_model(self, config: Mapping[str, Any]) -> object: - """Build a model through the user recipe.""" - if isinstance(self._recipe, TrainingRecipeV2): - modules = self._recipe.build_modules(config) - if not isinstance(modules, Mapping) or "model" not in modules: - raise AlgorithmConfigurationError( - "TrainingRecipeV2 build_modules must provide model" - ) - return modules["model"] - return self._recipe.model_factory(_mapping(config.get("model", {}), "model")) - - def build_optimizer(self, model: object, config: Mapping[str, Any]) -> object: - """Build an optimizer through the user recipe.""" - if isinstance(self._recipe, TrainingRecipeV2): - plan = self._recipe.optimization_plan( - model, - _mapping(config.get("optimizer", {}), "optimizer"), - ) - if not isinstance(plan, OptimizationPlan): - raise AlgorithmConfigurationError( - "TrainingRecipeV2 optimization_plan must return OptimizationPlan" - ) - return plan.optimizer - return self._recipe.optimizer_factory( - model, - _mapping(config.get("optimizer", {}), "optimizer"), - ) - - def build_loss(self, config: Mapping[str, Any]) -> object: - """Build a loss through the user recipe.""" - if isinstance(self._recipe, TrainingRecipeV2): - modules = self._recipe.build_modules(config) - if not isinstance(modules, Mapping) or "loss" not in modules: - raise AlgorithmConfigurationError( - "TrainingRecipeV2 build_modules must provide loss" - ) - return modules["loss"] - return self._recipe.loss_factory(_mapping(config.get("loss", {}), "loss")) - - def checkpoint_state(self, model: object, optimizer: object) -> Mapping[str, Any]: - """Return bounded replicated state for conformance tooling.""" - from tributo.training.distributed_torch import unwrapped_model - - model_state = getattr(unwrapped_model(model), "state_dict", None) - optimizer_state = getattr(optimizer, "state_dict", None) - if not callable(model_state) or not callable(optimizer_state): - raise AlgorithmConfigurationError( - "Torch recipe model and optimizer must expose state_dict" - ) - return {"model": model_state(), "optimizer": optimizer_state()} - - def train_loop_per_worker(self, config: Mapping[str, Any]) -> None: - """Run the framework-owned loop with the recipe's four factories.""" - worker_config = dict(config) - if isinstance(self._recipe, _FirstPartyTorchRecipeAdapter): - lowered = self._recipe._lower_worker_config(worker_config) - if not isinstance(lowered, Mapping): - raise AlgorithmConfigurationError( - "Torch recipe config adapter must return a mapping" - ) - worker_config = dict(lowered) - worker_config["_tributo_recipe_ref"] = str( - self._plan.implementation.implementation_ref - ) - worker_config["_tributo_recipe_code_digest"] = ( - self._plan.implementation.code_digest - ) - worker_config["_tributo_implementation_id"] = ( - self._plan.implementation.implementation_id - ) - worker_config["_tributo_algorithm"] = self._plan.resolution.algorithm - worker_config["_tributo_input_binding_digest"] = ( - self._plan.primary_input_descriptor.binding_digest - ) - worker_config["_tributo_distribution_spec_digest"] = ( - self._plan.runtime.distribution_digest - ) - worker_config["_tributo_resume_from"] = self._plan.runtime.resume_from - worker_config["_tributo_feature_names"] = list( - self._plan.primary_input_binding.feature_names - ) - worker_config["_tributo_label_name"] = ( - self._plan.primary_input_binding.label_name - ) - worker_config["_tributo_weight_name"] = ( - self._plan.primary_input_binding.sample_weight_name - ) - torch_recipe_train_loop_per_worker(worker_config, self._recipe) - - -def _batch_rows(value: object) -> int: - length = getattr(value, "__len__", None) - if not callable(length): - raise AlgorithmConfigurationError( - "Torch recipe batch values must expose a batch dimension" - ) - try: - return int(length()) - except (TypeError, ValueError) as exc: - raise AlgorithmConfigurationError( - "Torch recipe batch dimension must be an integer" - ) from exc - - -def _dense_batch( - batch: object, - *, - feature_names: tuple[str, ...], - label_name: str, - weight_name: str | None, -) -> tuple[Any, Any, Any | None, int]: - import torch - - if not isinstance(batch, Mapping): - raise AlgorithmConfigurationError("Ray Data Torch batches must be mappings") - required = (*feature_names, label_name, *((weight_name,) if weight_name else ())) - missing = [name for name in required if name not in batch] - if missing: - raise AlgorithmConfigurationError( - f"Torch recipe batch is missing required column(s): {missing}" - ) - labels = batch[label_name] - if not isinstance(labels, torch.Tensor): - raise AlgorithmConfigurationError("Ray Data label batches must be tensors") - rows = _batch_rows(labels) - columns: list[Any] = [] - for name in feature_names: - value = batch[name] - if not isinstance(value, torch.Tensor) or _batch_rows(value) != rows: - raise AlgorithmConfigurationError( - f"feature {name!r} must be a Tensor with the shared batch dimension" - ) - column = value.reshape(rows, -1).to(dtype=torch.float32) - if column.shape[1] != 1: - raise AlgorithmConfigurationError( - "the first Torch recipe input profile requires scalar feature columns" - ) - columns.append(column) - features = torch.cat(columns, dim=1) - targets = labels.to(dtype=torch.float32) - weights = None - if weight_name is not None: - weights = batch[weight_name] - if not isinstance(weights, torch.Tensor) or _batch_rows(weights) != rows: - raise AlgorithmConfigurationError( - "sample weights must be a Tensor with the shared batch dimension" - ) - weights = weights.to(dtype=torch.float64).reshape(-1) - if not bool(torch.isfinite(weights).all()) or bool((weights < 0).any()): - raise AlgorithmConfigurationError( - "sample weights must contain only finite non-negative values" - ) - return features, targets, weights, rows - - -def _prepare_batch( - recipe: TorchTrainingRecipe | TrainingRecipeV2, - batch: object, - *, - feature_names: tuple[str, ...], - label_name: str | None, - weight_name: str | None, - config: Mapping[str, Any], -) -> tuple[Any, Any, Any | None, int]: - if isinstance(recipe, TrainingRecipeV2): - prepared = recipe.batch_adapter( - batch, - feature_names=feature_names, - label_name=label_name, - weight_name=weight_name, - config=config, - ) - if not isinstance(prepared, tuple) or len(prepared) != 4: - raise AlgorithmConfigurationError( - "TrainingRecipeV2 batch_adapter must return four values" - ) - features, targets, weights, rows = prepared - if not isinstance(rows, int) or isinstance(rows, bool) or rows < 0: - raise AlgorithmConfigurationError( - "TrainingRecipeV2 batch_adapter returned an invalid row count" - ) - return features, targets, weights, rows - if label_name is None: - raise AlgorithmConfigurationError( - "TorchTrainingRecipe requires one label column" - ) - if not isinstance(recipe, _FirstPartyTorchRecipeAdapter): - return _dense_batch( - batch, - feature_names=feature_names, - label_name=label_name, - weight_name=weight_name, - ) - prepared = recipe._prepare_batch( - batch, - feature_names=feature_names, - label_name=label_name, - weight_name=weight_name, - config=config, - ) - if not isinstance(prepared, tuple) or len(prepared) != 4: - raise AlgorithmConfigurationError( - "Torch recipe batch adapter must return features, targets, weights, rows" - ) - features, targets, weights, rows = prepared - if not isinstance(rows, int) or isinstance(rows, bool) or rows < 0: - raise AlgorithmConfigurationError( - "Torch recipe batch adapter returned an invalid row count" - ) - return features, targets, weights, rows - - -def _empty_batch( - template: tuple[Any, Any, Any | None, int], -) -> tuple[Any, Any, Any | None, int]: - features, targets, weights, _ = template - empty_features = ( - {name: value[:0] for name, value in features.items()} - if isinstance(features, Mapping) - else features[:0] - ) - return ( - empty_features, - targets[:0], - weights[:0] if weights is not None else None, - 0, - ) - - -def _evaluation_metric_name(split: str, name: str) -> str: - if name.startswith("train_"): - return f"{split}_{name.removeprefix('train_')}" - return f"{split}_{name}" - - -def _metric_update( - state: dict[str, float], - value: object, - *, - reduction: MetricReduction, - rows: int, - weights: Any | None, -) -> None: - import torch - - tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) - tensor = tensor.detach().to(dtype=torch.float64).reshape(-1) - if tensor.numel() == 0 or not bool(torch.isfinite(tensor).all()): - raise AlgorithmConfigurationError( - "Torch recipe metric produced empty or non-finite values" - ) - if reduction is MetricReduction.SUM_COUNT: - if tensor.numel() == 1: - state["value"] += float(tensor.item()) * rows - state["weight"] += rows - else: - state["value"] += float(tensor.sum().item()) - state["weight"] += int(tensor.numel()) - return - if reduction is MetricReduction.WEIGHTED_MEAN: - if weights is None: - raise AlgorithmConfigurationError( - "weighted_mean metric requires an InputBinding sample_weight_name" - ) - metric_values = tensor - if metric_values.numel() == 1: - metric_values = metric_values.expand_as(weights) - if metric_values.numel() != weights.numel(): - raise AlgorithmConfigurationError( - "weighted metric values must match the sample-weight tensor" - ) - state["value"] += float((metric_values * weights).sum().item()) - state["weight"] += float(weights.sum().item()) - return - candidate = float( - (tensor.min() if reduction is MetricReduction.MIN else tensor.max()).item() - ) - if reduction is MetricReduction.MIN: - state["value"] = min(state["value"], candidate) - else: - state["value"] = max(state["value"], candidate) - - -def _reduce_metrics( - states: Mapping[str, dict[str, float]], - reducers: Mapping[str, MetricReduction], -) -> dict[str, float]: - from tributo.training.distributed_torch import ( - all_gather_objects, - all_reduce_values, - ) - - reduced: dict[str, float] = {} - for name in sorted(reducers): - reduction = reducers[name] - state = states[name] - if reduction in {MetricReduction.SUM_COUNT, MetricReduction.WEIGHTED_MEAN}: - value, weight = all_reduce_values((state["value"], state["weight"])) - if weight <= 0: - raise AlgorithmConfigurationError( - f"metric {name!r} has a non-positive global weight" - ) - reduced[name] = value / weight - else: - candidates = tuple( - float(value) for value in all_gather_objects(state["value"]) - ) - finite = [value for value in candidates if math.isfinite(value)] - if not finite: - raise AlgorithmConfigurationError( - f"metric {name!r} has no finite global value" - ) - reduced[name] = ( - min(finite) if reduction is MetricReduction.MIN else max(finite) - ) - return reduced - - -def _checkpoint_contract( - *, - config: Mapping[str, Any], - feature_count: int, - output_shape: tuple[int, ...], - framework_version: str, -) -> dict[str, Any]: - from tributo.exporting.models import CheckpointField, ExportCheckpointV1 - - feature_names = tuple(config.get("_tributo_feature_names") or ()) - if len(feature_names) != feature_count or any( - not isinstance(name, str) or not name for name in feature_names - ): - raise AlgorithmConfigurationError( - "Torch recipe checkpoint feature declaration is invalid" - ) - contract = ExportCheckpointV1( - trainer_type=_TRAINER_TYPE, - architecture_id=str(config["_tributo_implementation_id"]), - input_schema=tuple( - CheckpointField( - name=name, - dtype="float32", - shape=("batch",), - ) - for name in feature_names - ), - output_schema=( - CheckpointField( - name="output", - dtype="float32", - shape=("batch", *output_shape), - ), - ), - task_type=str(config["_tributo_algorithm"]), - framework="pytorch", - framework_version=framework_version, - required_artifacts=("model.pt",), - ).model_dump(mode="json") - contract.update( - { - "model": dict(_mapping(config.get("model", {}), "model")), - "recipe_ref": str(config["_tributo_recipe_ref"]), - "recipe_code_digest": config.get("_tributo_recipe_code_digest"), - } - ) - return contract - - -def _recipe_checkpoint_contract( - recipe: TorchTrainingRecipe | TrainingRecipeV2, - *, - config: Mapping[str, Any], - feature_count: int, - output_shape: tuple[int, ...], - framework_version: str, - model_digest: str, - world_size: int, -) -> dict[str, Any]: - if not isinstance(recipe, _FirstPartyTorchRecipeAdapter): - return _checkpoint_contract( - config=config, - feature_count=feature_count, - output_shape=output_shape, - framework_version=framework_version, - ) - contract = recipe._checkpoint_contract( - config=config, - feature_count=feature_count, - output_shape=output_shape, - framework_version=framework_version, - model_digest=model_digest, - world_size=world_size, - ) - if not isinstance(contract, dict): - raise AlgorithmConfigurationError( - "Torch recipe checkpoint adapter must return a dictionary" - ) - return contract - - -def _write_recipe_checkpoint_artifacts( - recipe: TorchTrainingRecipe | TrainingRecipeV2, - checkpoint_dir: Path, -) -> tuple[str, ...]: - if not isinstance(recipe, _FirstPartyTorchRecipeAdapter): - return () - files = recipe._write_checkpoint_artifacts(checkpoint_dir) - if not isinstance(files, tuple) or any( - not isinstance(name, str) - or not name - or Path(name).name != name - or not (checkpoint_dir / name).is_file() - for name in files - ): - raise AlgorithmConfigurationError( - "Torch recipe checkpoint adapter returned invalid artifact files" - ) - return files - - -def _validate_resume( - metadata: Mapping[str, Any], - *, - world_size: int, - distribution_digest: str | None, -) -> None: - if metadata.get("world_size") != world_size: - raise AlgorithmConfigurationError( - "Torch recipe resume requires the original world size" - ) - if metadata.get("distribution_spec_digest") != distribution_digest: - raise AlgorithmConfigurationError( - "Torch recipe resume DistributionSpec digest does not match" - ) - - -def _evaluate_dataset( - data: object, - *, - split: str, - recipe: TorchTrainingRecipe | TrainingRecipeV2, - model: object, - loss_fn: object, - metrics: Mapping[str, Any], - reducers: Mapping[str, MetricReduction], - loop: _LoopConfig, - feature_names: tuple[str, ...], - label_name: str | None, - weight_name: str | None, - modules: Mapping[str, object], - config: Mapping[str, Any], -) -> tuple[dict[str, float], int]: - import torch - - from tributo.training.distributed_torch import all_reduce_values - - iterate = getattr(data, "iter_torch_batches", None) - if not callable(iterate): - raise AlgorithmConfigurationError( - f"Torch recipe {split} input is not a Ray Data iterator" - ) - iterator = iter( - iterate( - batch_size=loop.batch_size, - prefetch_batches=loop.prefetch_batches, - dtypes=torch.float32, - drop_last=False, - ) - ) - states: dict[str, dict[str, float]] = { - name: { - "value": ( - float("inf") - if reduction is MetricReduction.MIN - else float("-inf") - if reduction is MetricReduction.MAX - else 0.0 - ), - "weight": 0.0, - } - for name, reduction in reducers.items() - } - local_rows = 0 - unwrapped = getattr(model, "module", model) - train = getattr(unwrapped, "train", None) - if not callable(train): - raise AlgorithmConfigurationError("Torch recipe model must expose train()") - train(False) - current = next(iterator, None) - while True: - rows = 0 - prepared = None - if current is not None: - prepared = _prepare_batch( - recipe, - current, - feature_names=feature_names, - label_name=label_name, - weight_name=weight_name, - config=config, - ) - rows = prepared[3] - (global_rows,) = all_reduce_values((float(rows),)) - if global_rows <= 0: - break - if prepared is not None: - features, targets, weights, rows = prepared - with torch.no_grad(): - if isinstance(recipe, TrainingRecipeV2): - step = recipe.validation_step( - {**modules, "model": unwrapped}, - features, - targets, - weights, - config, - ) - if not isinstance(step, TrainingStepResult): - raise AlgorithmConfigurationError( - "TrainingRecipeV2 validation_step must return " - "TrainingStepResult" - ) - predictions = step.predictions - loss = step.loss - else: - predictions = recipe.forward(unwrapped, features) - if not isinstance(predictions, torch.Tensor): - raise AlgorithmConfigurationError( - "Torch recipe forward must return one Tensor" - ) - aligned_targets = targets - if predictions.numel() == targets.numel(): - aligned_targets = targets.reshape_as(predictions) - if not isinstance(recipe, TrainingRecipeV2): - loss = recipe.compute_loss(loss_fn, predictions, aligned_targets) - if not isinstance(loss, torch.Tensor) or loss.ndim != 0: - raise AlgorithmConfigurationError( - "Torch recipe evaluation loss must return one scalar batch mean" - ) - _metric_update( - states["train_loss"], - loss, - reduction=MetricReduction.SUM_COUNT, - rows=rows, - weights=None, - ) - for name, metric in metrics.items(): - _metric_update( - states[name], - metric(predictions, aligned_targets), - reduction=reducers[name], - rows=rows, - weights=weights, - ) - local_rows += rows - current = next(iterator, None) - reduced = _reduce_metrics(states, reducers) - return { - f"{split}_loss": reduced["train_loss"], - **{ - _evaluation_metric_name(split, name): value - for name, value in reduced.items() - if name != "train_loss" - }, - }, local_rows - - -@DeveloperAPI -def torch_recipe_train_loop_per_worker( - config: Mapping[str, Any], - recipe: TorchTrainingRecipe | TrainingRecipeV2, -) -> None: - """Execute one dense-tabular recipe with Ray Data and replicated DDP.""" - import numpy as np - import ray.train - import torch - import torch.distributed as dist - from ray.train.torch import enable_reproducibility - - from tributo.training.checkpoint import ( - ResumeConfig, - capture_rng_state, - checkpoint_directory, - read_resume_manifest, - restore_rng_state, - write_resume_manifest, - ) - from tributo.training.distributed_torch import ( - all_gather_objects, - all_reduce_values, - broadcast_bool, - collective_execution_evidence, - prepare_model, - unwrapped_model, - ) - - loop = _LoopConfig.model_validate(config.get("training") or {}) - ray_config = _mapping(config.get("ray", {}), "ray config") - resume = ResumeConfig.model_validate(ray_config.get("resume") or {}) - feature_names = tuple(config.get("_tributo_feature_names") or ()) - label_name = config.get("_tributo_label_name") - weight_name = config.get("_tributo_weight_name") - if not feature_names: - raise AlgorithmConfigurationError("Torch recipe requires feature columns") - if not isinstance(recipe, TrainingRecipeV2) and ( - not isinstance(label_name, str) or not label_name - ): - raise AlgorithmConfigurationError( - "TorchTrainingRecipe requires one label column" - ) - if label_name is not None and (not isinstance(label_name, str) or not label_name): - raise AlgorithmConfigurationError( - "TrainingRecipeV2 label column must be non-empty when provided" - ) - if weight_name is not None and not isinstance(weight_name, str): - raise AlgorithmConfigurationError("sample weight column must be a string") - - context = ray.train.get_context() - rank = context.get_world_rank() - world_size = context.get_world_size() - trainer_type = ( - recipe._trainer_type - if isinstance(recipe, _FirstPartyTorchRecipeAdapter) - else _TRAINER_TYPE - ) - enable_reproducibility(loop.seed) - random.seed(loop.seed + rank) - np.random.seed(loop.seed + rank) - - modules: dict[str, object] - optimization_plan: OptimizationPlan | None = None - if isinstance(recipe, TrainingRecipeV2): - built_modules = recipe.build_modules(config) - if ( - not isinstance(built_modules, Mapping) - or "model" not in built_modules - or "loss" not in built_modules - ): - raise AlgorithmConfigurationError( - "TrainingRecipeV2 build_modules must provide model and loss" - ) - modules = dict(built_modules) - model: Any = modules["model"] - loss_fn = modules["loss"] - optimization_plan = recipe.optimization_plan( - model, - _mapping(config.get("optimizer", {}), "optimizer"), - ) - if not isinstance(optimization_plan, OptimizationPlan): - raise AlgorithmConfigurationError( - "TrainingRecipeV2 optimization_plan must return OptimizationPlan" - ) - optimizer: Any = optimization_plan.optimizer - metric_plan = recipe.metric_plan(_mapping(config.get("metrics", {}), "metrics")) - if not isinstance(metric_plan, MetricPlan): - raise AlgorithmConfigurationError( - "TrainingRecipeV2 metric_plan must return MetricPlan" - ) - metrics = metric_plan.factories - else: - model = recipe.model_factory(_mapping(config.get("model", {}), "model")) - loss_fn = recipe.loss_factory(_mapping(config.get("loss", {}), "loss")) - optimizer = recipe.optimizer_factory( - model, - _mapping(config.get("optimizer", {}), "optimizer"), - ) - metrics = recipe.metric_factories( - _mapping(config.get("metrics", {}), "metrics") - ) - modules = {"model": model, "loss": loss_fn} - if not isinstance(model, torch.nn.Module): - raise AlgorithmConfigurationError( - "Torch recipe model_factory must return torch.nn.Module" - ) - if world_size > 1 and any( - isinstance(module, torch.nn.modules.batchnorm._BatchNorm) - for module in model.modules() - ): - raise AlgorithmConfigurationError( - "replicated Torch recipes reject BatchNorm until synchronized buffers " - "have a separate gate" - ) - if not callable(loss_fn): - raise AlgorithmConfigurationError( - "Torch recipe loss_factory must return a callable" - ) - for method in ("zero_grad", "step", "state_dict", "load_state_dict"): - if not callable(getattr(optimizer, method, None)): - raise AlgorithmConfigurationError( - "Torch recipe optimizer does not implement the optimizer contract" - ) - if not isinstance(metrics, Mapping) or any( - not isinstance(name, str) or not name or not callable(metric) - for name, metric in metrics.items() - ): - raise AlgorithmConfigurationError( - "Torch recipe metric_factories must return named callables" - ) - reducer_payload = _mapping( - config.get("_tributo_metric_reducers", {}), - "metric reducers", - ) - try: - reducers = { - name: MetricReduction(value) for name, value in reducer_payload.items() - } - except (TypeError, ValueError) as exc: - raise AlgorithmConfigurationError("Torch metric reducer is invalid") from exc - expected_metric_names = set(reducers) - {"train_loss"} - if set(metrics) != expected_metric_names: - raise AlgorithmConfigurationError( - "Torch recipe metric names must exactly match CollectivePolicy reducers" - ) - scheduler: Any = ( - optimization_plan.scheduler if optimization_plan is not None else None - ) - if scheduler is not None and any( - not callable(getattr(scheduler, method, None)) - for method in ("step", "state_dict", "load_state_dict") - ): - raise AlgorithmConfigurationError( - "TrainingRecipeV2 scheduler does not implement the scheduler contract" - ) - scaler = torch.amp.GradScaler("cuda", enabled=loop.amp) - - checkpoint = ray.train.get_checkpoint() - if checkpoint is None: - explicit_checkpoint = config.get("_tributo_resume_from") - if explicit_checkpoint is not None and not isinstance(explicit_checkpoint, str): - raise AlgorithmConfigurationError( - "Torch recipe explicit checkpoint path must be a string" - ) - checkpoint = explicit_checkpoint - start_epoch = 0 - best_val_loss = float("inf") - patience_counter = 0 - if checkpoint is not None: - with checkpoint_directory(checkpoint) as checkpoint_dir: - envelope = read_resume_manifest( - checkpoint_dir, - expected_trainer_type=trainer_type, - expected_resume_id=resume.resume_id, - ) - _validate_resume( - envelope.payload_metadata, - world_size=world_size, - distribution_digest=config.get("_tributo_distribution_spec_digest"), - ) - model.load_state_dict( - torch.load( - checkpoint_dir / "model.pt", - map_location="cpu", - weights_only=True, - ) - ) - optimizer.load_state_dict( - torch.load( - checkpoint_dir / "optimizer.pt", - map_location="cpu", - weights_only=True, - ) - ) - scheduler_path = checkpoint_dir / "scheduler.pt" - if scheduler is not None: - if not scheduler_path.is_file(): - raise AlgorithmConfigurationError( - "TrainingRecipeV2 checkpoint is missing scheduler state" - ) - scheduler.load_state_dict( - torch.load(scheduler_path, map_location="cpu", weights_only=True) - ) - scaler_path = checkpoint_dir / "scaler.pt" - if scaler_path.is_file(): - scaler.load_state_dict( - torch.load(scaler_path, map_location="cpu", weights_only=True) - ) - rng = json.loads((checkpoint_dir / "rng_state.json").read_text()) - training_state_path = checkpoint_dir / "training_state.json" - training_state = ( - json.loads(training_state_path.read_text()) - if training_state_path.is_file() - else {} - ) - if isinstance(recipe, _FirstPartyTorchRecipeAdapter): - recipe._validate_checkpoint_artifacts(checkpoint_dir) - rank_states = rng.get("rank_states") if isinstance(rng, dict) else None - if not isinstance(rank_states, list) or len(rank_states) != world_size: - raise AlgorithmConfigurationError( - "Torch recipe checkpoint RNG state does not match world size" - ) - restore_rng_state(rank_states[rank]) - start_epoch = envelope.completed_step - best_val_loss = float(training_state.get("best_val_loss", float("inf"))) - patience_counter = int(training_state.get("patience_counter", 0)) - - model, device = prepare_model(model) - modules["model"] = model - if loop.amp and getattr(device, "type", None) != "cuda": - raise AlgorithmConfigurationError("Torch recipe AMP requires a CUDA worker") - train_data = ray.train.get_dataset_shard("train") - if train_data is None: - raise AlgorithmConfigurationError("Torch recipe did not receive train data") - evaluation_data: dict[str, object] = {} - for split in ("val", "test"): - try: - shard = ray.train.get_dataset_shard(split) - except KeyError: - shard = None - if shard is not None: - evaluation_data[split] = shard - - output_shape: tuple[int, ...] | None = None - for epoch in range(start_epoch, loop.epochs): - iterator = iter( - train_data.iter_torch_batches( - batch_size=loop.batch_size, - prefetch_batches=loop.prefetch_batches, - dtypes=torch.float32, - drop_last=False, - local_shuffle_buffer_size=loop.local_shuffle_buffer_size, - local_shuffle_seed=loop.seed + epoch, - ) - ) - first_raw = next(iterator, None) - local_non_empty = 1.0 if first_raw is not None else 0.0 - (non_empty_ranks,) = all_reduce_values((local_non_empty,)) - if int(non_empty_ranks) != world_size: - raise AlgorithmConfigurationError( - "exact-coverage Torch recipes require at least one batch per rank" - ) - assert first_raw is not None - template = _prepare_batch( - recipe, - first_raw, - feature_names=feature_names, - label_name=label_name, - weight_name=weight_name, - config=config, - ) - current: tuple[Any, Any, Any | None, int] | None = template - states: dict[str, dict[str, float]] = { - name: { - "value": ( - float("inf") - if reduction is MetricReduction.MIN - else float("-inf") - if reduction is MetricReduction.MAX - else 0.0 - ), - "weight": 0.0, - } - for name, reduction in reducers.items() - } - local_rows = 0 - local_batches = 0 - collective_steps = 0 - local_coverage_counts: dict[str, int] = {} - model.train() - optimizer.zero_grad() - while True: - active = current is not None - prepared = _empty_batch(template) if current is None else current - features, targets, weights, rows = prepared - (global_rows,) = all_reduce_values((float(rows),)) - if global_rows <= 0: - break - with torch.autocast( - device_type=device.type, - enabled=loop.amp, - ): - if isinstance(recipe, TrainingRecipeV2): - step = recipe.training_step( - modules, - features, - targets, - weights, - config, - ) - if not isinstance(step, TrainingStepResult): - raise AlgorithmConfigurationError( - "TrainingRecipeV2 training_step must return " - "TrainingStepResult" - ) - predictions = step.predictions - raw_loss = step.loss - if active: - for name, count in step.coverage_counts.items(): - local_coverage_counts[name] = ( - local_coverage_counts.get(name, 0) + count - ) - else: - predictions = recipe.forward(model, features) - if not isinstance(predictions, torch.Tensor): - raise AlgorithmConfigurationError( - "Torch recipe forward must return one Tensor" - ) - if output_shape is None and active: - output_shape = tuple(int(size) for size in predictions.shape[1:]) - aligned_targets = targets - if predictions.numel() == targets.numel(): - aligned_targets = targets.reshape_as(predictions) - if active and not isinstance(recipe, TrainingRecipeV2): - raw_loss = recipe.compute_loss( - loss_fn, - predictions, - aligned_targets, - ) - if not isinstance(raw_loss, torch.Tensor) or raw_loss.ndim != 0: - raise AlgorithmConfigurationError( - "Torch recipe loss must return one scalar batch mean" - ) - if not bool(torch.isfinite(raw_loss)): - raise AlgorithmConfigurationError( - "Torch recipe loss produced a non-finite value" - ) - local_loss_sum = raw_loss * rows - elif not active: - raw_loss = predictions.sum() * 0.0 - local_loss_sum = raw_loss - else: - if not isinstance(raw_loss, torch.Tensor) or raw_loss.ndim != 0: - raise AlgorithmConfigurationError( - "TrainingRecipeV2 loss must return one scalar batch mean" - ) - if not bool(torch.isfinite(raw_loss)): - raise AlgorithmConfigurationError( - "TrainingRecipeV2 loss produced a non-finite value" - ) - local_loss_sum = raw_loss * rows - backward_loss = local_loss_sum * world_size / global_rows - scaler.scale(backward_loss).backward() - accumulation = ( - optimization_plan.gradient_accumulation_steps - if optimization_plan is not None - else 1 - ) - if (collective_steps + 1) % accumulation == 0: - if ( - optimization_plan is not None - and optimization_plan.max_gradient_norm is not None - ): - scaler.unscale_(optimizer) - torch.nn.utils.clip_grad_norm_( - model.parameters(), - optimization_plan.max_gradient_norm, - ) - scaler.step(optimizer) - scaler.update() - optimizer.zero_grad() - - if active: - _metric_update( - states["train_loss"], - raw_loss, - reduction=MetricReduction.SUM_COUNT, - rows=rows, - weights=None, - ) - for name, metric in metrics.items(): - value = metric(predictions.detach(), aligned_targets.detach()) - _metric_update( - states[name], - value, - reduction=reducers[name], - rows=rows, - weights=weights, - ) - local_rows += rows - local_batches += 1 - collective_steps += 1 - raw = next(iterator, None) - current = ( - _prepare_batch( - recipe, - raw, - feature_names=feature_names, - label_name=label_name, - weight_name=weight_name, - config=config, - ) - if raw is not None - else None - ) - - accumulation = ( - optimization_plan.gradient_accumulation_steps - if optimization_plan is not None - else 1 - ) - remainder = collective_steps % accumulation - if remainder: - gradient_scale = accumulation / remainder - for parameter in model.parameters(): - if parameter.grad is not None: - parameter.grad.mul_(gradient_scale) - if ( - optimization_plan is not None - and optimization_plan.max_gradient_norm is not None - ): - scaler.unscale_(optimizer) - torch.nn.utils.clip_grad_norm_( - model.parameters(), - optimization_plan.max_gradient_norm, - ) - scaler.step(optimizer) - scaler.update() - optimizer.zero_grad() - - epoch_metrics = {"epoch": epoch + 1, **_reduce_metrics(states, reducers)} - if scheduler is not None: - scheduler.step() - input_rows = {"train": local_rows} - input_rows.update( - { - f"coverage.{name}": count - for name, count in sorted(local_coverage_counts.items()) - } - ) - for split, data in evaluation_data.items(): - if split == "test" and epoch + 1 != loop.epochs: - continue - split_metrics, split_rows = _evaluate_dataset( - data, - split=split, - recipe=recipe, - model=model, - loss_fn=loss_fn, - metrics=metrics, - reducers=reducers, - loop=loop, - feature_names=feature_names, - label_name=label_name, - weight_name=weight_name, - modules=modules, - config=config, - ) - epoch_metrics.update(split_metrics) - input_rows[split] = split_rows - if any(not math.isfinite(float(value)) for value in epoch_metrics.values()): - raise AlgorithmConfigurationError( - "Torch recipe produced non-finite global metrics" - ) - stop_after_report = False - if loop.early_stopping_patience is not None and "val_loss" in epoch_metrics: - if rank == 0: - val_loss = float(epoch_metrics["val_loss"]) - if val_loss < best_val_loss: - best_val_loss = val_loss - patience_counter = 0 - else: - patience_counter += 1 - stop_after_report = patience_counter >= loop.early_stopping_patience - stop_after_report = broadcast_bool(stop_after_report, source_rank=0) - execution_workers, model_digest = collective_execution_evidence( - model, - shard_rows=local_rows, - input_binding_digest=config.get("_tributo_input_binding_digest"), - input_rows=input_rows, - batch_count=local_batches, - collective_steps=collective_steps, - ) - report: dict[str, Any] = { - **epoch_metrics, - "execution_workers": list(execution_workers), - "model_state_digest": model_digest, - "world_size": world_size, - "state_coordination": "all_reduce", - "collective_backend": ( - str(dist.get_backend()) if dist.is_initialized() else "none" - ), - "checkpoint_owner_rank": 0, - "metric_reducers": { - name: reducer.value for name, reducer in reducers.items() - }, - } - should_checkpoint = ( - (epoch + 1) % resume.checkpoint_interval == 0 - or epoch + 1 == loop.epochs - or stop_after_report - ) - rank_rng_states = ( - all_gather_objects(capture_rng_state()) if should_checkpoint else () - ) - if rank == 0 and should_checkpoint: - from ray.train import Checkpoint - - checkpoint_dir = Path(tempfile.mkdtemp(prefix="torch_recipe_ckpt_")) - try: - torch.save( - unwrapped_model(model).state_dict(), - checkpoint_dir / "model.pt", - ) - torch.save(optimizer.state_dict(), checkpoint_dir / "optimizer.pt") - torch.save(scaler.state_dict(), checkpoint_dir / "scaler.pt") - scheduler_files: tuple[str, ...] = () - if scheduler is not None: - torch.save( - scheduler.state_dict(), - checkpoint_dir / "scheduler.pt", - ) - scheduler_files = ("scheduler.pt",) - if output_shape is None: - raise AlgorithmConfigurationError( - "Torch recipe did not observe a model output shape" - ) - model_config = _recipe_checkpoint_contract( - recipe, - config=config, - feature_count=len(feature_names), - output_shape=output_shape, - framework_version=torch.__version__, - model_digest=model_digest, - world_size=world_size, - ) - (checkpoint_dir / "model_config.json").write_text( - json.dumps(model_config, ensure_ascii=False), - encoding="utf-8", - ) - (checkpoint_dir / "metrics.json").write_text( - json.dumps(epoch_metrics, ensure_ascii=False), - encoding="utf-8", - ) - (checkpoint_dir / "rng_state.json").write_text( - json.dumps({"rank_states": list(rank_rng_states)}), - encoding="utf-8", - ) - (checkpoint_dir / "training_state.json").write_text( - json.dumps( - { - "best_val_loss": best_val_loss, - "patience_counter": patience_counter, - } - ), - encoding="utf-8", - ) - extra_files = _write_recipe_checkpoint_artifacts( - recipe, - checkpoint_dir, - ) - envelope = write_resume_manifest( - checkpoint_dir, - resume_id=resume.resume_id, - trainer_type=trainer_type, - completed_step=epoch + 1, - framework="pytorch", - framework_version=torch.__version__, - payload_files=( - "metrics.json", - "model.pt", - "model_config.json", - "optimizer.pt", - "scaler.pt", - *scheduler_files, - "rng_state.json", - "training_state.json", - *extra_files, - ), - payload_metadata={ - "world_size": world_size, - "distribution_spec_digest": config.get( - "_tributo_distribution_spec_digest" - ), - **( - {"preprocessing": "preprocessor.json"} - if "preprocessor.json" in extra_files - else {} - ), - }, - ) - report["resume_id"] = envelope.resume_id - ray.train.report( - report, - checkpoint=Checkpoint.from_directory(str(checkpoint_dir)), - ) - finally: - shutil.rmtree(checkpoint_dir, ignore_errors=True) - else: - ray.train.report(report) - if stop_after_report: - break - - -@DeveloperAPI -def create_torch_recipe_algorithm( - *, - plan: ResolvedAlgorithmPlan, - implementation: object, - artifacts: tuple[object, ...], -) -> CollectiveAlgorithm: - """Construct the internal collective adapter for one trusted recipe class.""" - del artifacts - recipe_cls = _recipe_type( - plan.implementation.implementation_ref, - plan.implementation.code_digest, - ) - if implementation is not recipe_cls: - raise AlgorithmConfigurationError( - "Torch recipe implementation drifted after descriptor resolution" - ) - return _TorchRecipeCollectiveAlgorithm( - plan, - _new_recipe( - plan.implementation.implementation_ref, - plan.implementation.code_digest, - ), - ) - - -@DeveloperAPI -def export_torch_recipe_result( - *, - result: object, - plan: ResolvedAlgorithmPlan, - run_id: str, -) -> AlgorithmExecutionResult: - """Publish a recipe checkpoint through the existing ONNX Bundle pipeline.""" - import importlib.metadata - - from tributo.exporting.models import BundleOutputConfig, ExportTarget - from tributo.exporting.service import BundleExportService - from tributo.integrations.algorithm_runtimes.portable_metrics import ( - portable_fit_only_metrics, - ) - from tributo.integrations.sources.ray_torch_recipe import ( - RayTorchRecipeSourceProvider, - TorchRecipeSourceOptions, - ) - - output = _OutputConfig.model_validate(plan.algorithm_config.get("output") or {}) - provider = RayTorchRecipeSourceProvider() - options = TorchRecipeSourceOptions( - recipe_ref=str(plan.implementation.implementation_ref), - recipe_code_digest=plan.implementation.code_digest, - implementation_id=plan.implementation.implementation_id, - ) - bundle_config = BundleOutputConfig( - bundle_uri=output.bundle_uri, - request_id=run_id, - run_id=run_id, - targets=[ - ExportTarget( - name="onnx-model", - format="onnx", - exporter_id="torch-onnx-v1", - options={"opset": 18}, - ) - ], - roles={"inference": "onnx-model"}, - ) - with provider.open_source(result, options) as source: - published = BundleExportService().export_bundle( - source, - bundle_config, - tributo_version=importlib.metadata.version("tributo"), - ) - raw_metrics = getattr(result, "metrics", None) or {} - return AlgorithmExecutionResult( - status="succeeded", - metrics=portable_fit_only_metrics(raw_metrics), - outputs={ - "bundle_id": published.bundle_id, - "bundle_uri": published.canonical_uri, - "execution_id": published.execution_id, - "manifest_sha256": published.manifest_sha256, - }, - ) - - -__all__ = [ - "create_torch_recipe_algorithm", - "export_torch_recipe_result", - "torch_recipe_train_loop_per_worker", -] diff --git a/src/tributo/integrations/sources/__init__.py b/src/tributo/integrations/sources/__init__.py index 8ccd4cc..60f1287 100644 --- a/src/tributo/integrations/sources/__init__.py +++ b/src/tributo/integrations/sources/__init__.py @@ -2,6 +2,6 @@ from __future__ import annotations -from tributo.integrations.sources.ray_torch_recipe import RayTorchRecipeSourceProvider +from tributo.integrations.sources.ray_torch import RayTorchSourceProvider -__all__ = ["RayTorchRecipeSourceProvider"] +__all__ = ["RayTorchSourceProvider"] diff --git a/src/tributo/integrations/sources/ray_torch.py b/src/tributo/integrations/sources/ray_torch.py new file mode 100644 index 0000000..81c9e42 --- /dev/null +++ b/src/tributo/integrations/sources/ray_torch.py @@ -0,0 +1,542 @@ +"""Source Provider for final-stage Core Torch checkpoints.""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from contextlib import contextmanager +from pathlib import Path +from typing import Any, ClassVar, Generator + +from pydantic import BaseModel, ConfigDict, Field + +from tributo.algorithms.api import ( + AlgorithmConfigurationError, + QualifiedReference, + TorchCheckpointDescriptor, + TorchCheckpointRef, +) +from tributo.algorithms.core.worker import _load_reference, _validate_module_digest +from tributo.algorithms.spi import ( + RayTorchAdapter, + TorchArtifactContext, + TorchArtifactPlan, + TorchRecipe, + TorchRuntimeContext, +) +from tributo.exporting.models import ( + CheckpointField, + ExportCheckpointV1, + ExportSource, +) +from tributo.training.checkpoint import checkpoint_directory +from tributo.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class TorchSourceOptions(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + implementation_ref: str = Field(min_length=1) + implementation_code_digest: str = Field(pattern=r"^[0-9a-f]{64}$") + implementation_id: str = Field(min_length=1) + loop_owner: str = Field(default="core_recipe", pattern=r"^(core_recipe|adapter)$") + algorithm_config: dict[str, Any] = Field(default_factory=dict) + input_bindings: dict[str, Any] = Field(default_factory=dict) + output_config: dict[str, Any] = Field(default_factory=dict) + policy_digest: str = Field(pattern=r"^[0-9a-f]{64}$") + plan_digest: str = Field(pattern=r"^[0-9a-f]{64}$") + input_binding_digest: str = Field(pattern=r"^[0-9a-f]{64}$") + stage_input_roles: tuple[str, ...] = ("train",) + stage_index: int = Field(default=0, ge=0) + + +def _source_context_config(options: TorchSourceOptions) -> dict[str, Any]: + """Keep Core Ray/output control fields out of implementation contexts.""" + return { + str(key): value + for key, value in options.algorithm_config.items() + if str(key) not in {"ray", "output"} + } + + +def _checkpoint_contract_from_artifact_plan( + artifact_payload: Mapping[str, Any], + descriptor: TorchCheckpointDescriptor, + options: TorchSourceOptions, + *, + preprocessing: Mapping[str, Any] | None = None, +) -> ExportCheckpointV1: + """Translate the typed Torch artifact declaration to Bundle metadata.""" + import torch + + try: + input_schema = tuple( + CheckpointField.model_validate(dict(field)) + for field in artifact_payload.get("input_signature", ()) + ) + output_schema = tuple( + CheckpointField.model_validate(dict(field)) + for field in artifact_payload.get("output_signature", ()) + ) + except (TypeError, ValueError) as exc: + raise AlgorithmConfigurationError( + "Torch artifact plan signature is malformed" + ) from exc + if not input_schema or not output_schema: + raise AlgorithmConfigurationError( + "Torch artifact plan requires input and output signatures" + ) + targets = artifact_payload.get("targets", ()) + required_artifacts = tuple( + str(target["name"]) + for target in targets + if isinstance(target, Mapping) and isinstance(target.get("name"), str) + ) + configured_task = options.algorithm_config.get("task_type") + task_type = ( + configured_task + if isinstance(configured_task, str) and configured_task + else descriptor.identity.algorithm + ) + return ExportCheckpointV1( + trainer_type="ray_train_torch", + architecture_id=descriptor.identity.implementation_id, + input_schema=input_schema, + output_schema=output_schema, + preprocessing=dict(preprocessing or {}), + task_type=task_type, + framework="pytorch", + framework_version=str(torch.__version__), + checkpoint_format_version=1, + required_artifacts=required_artifacts, + ) + + +@PublicAPI(stability="alpha") +class RayTorchSourceProvider: + """Open an ExportSource while preserving the checkpoint lease lifetime.""" + + api_version: ClassVar[int] = 1 + provider_id: ClassVar[str] = "ray-torch-v1" + trainer_type: ClassVar[str] = "ray_train_torch" + priority: ClassVar[int] = 100 + + def open_source(self, result: Any, config: BaseModel | None = None) -> Any: + options = TorchSourceOptions.model_validate( + config.model_dump() if config is not None else {} + ) + return self.open_export_source(result, options) + + def open_export_source(self, result: Any, config: TorchSourceOptions) -> Any: + return _open_source(result, config) + + +@contextmanager +def _open_source( + result: Any, + options: TorchSourceOptions, +) -> Generator[ExportSource, None, None]: + checkpoint = getattr(result, "checkpoint", result) + if checkpoint is None: + raise AlgorithmConfigurationError("Torch result has no final Checkpoint") + with checkpoint_directory(checkpoint) as checkpoint_dir: + descriptor = _read_descriptor(checkpoint_dir) + reference = options.implementation_ref.partition(":") + if not reference[1]: + raise AlgorithmConfigurationError( + "Torch implementation reference is invalid" + ) + qualified = f"{reference[0]}:{reference[2]}" + if ( + descriptor.adapter_identity is not None + and descriptor.adapter_identity != options.implementation_id + ): + raise AlgorithmConfigurationError( + "Torch checkpoint implementation identity drifted" + ) + if descriptor.identity.implementation_id != options.implementation_id: + raise AlgorithmConfigurationError( + "Torch checkpoint implementation identity drifted" + ) + if ( + descriptor.identity.implementation_code_digest + != options.implementation_code_digest + ): + raise AlgorithmConfigurationError( + "Torch checkpoint implementation code digest drifted" + ) + implementation_ref = QualifiedReference.parse(qualified) + _validate_module_digest(implementation_ref, options.implementation_code_digest) + implementation = _load_reference(implementation_ref) + if options.policy_digest != descriptor.policy_digest: + raise AlgorithmConfigurationError( + "Torch export Policy digest does not match checkpoint" + ) + if descriptor.identity.plan_digest != options.plan_digest: + raise AlgorithmConfigurationError( + "Torch export plan digest does not match checkpoint" + ) + if not isinstance(implementation, type) or not issubclass( + implementation, (TorchRecipe, RayTorchAdapter) + ): + raise AlgorithmConfigurationError( + "Torch export implementation is not a Recipe or Adapter" + ) + checkpoint_ref = TorchCheckpointRef( + checkpoint=checkpoint, + descriptor_digest=descriptor.digest, + source_stage_id=descriptor.identity.stage_id, + descriptor=descriptor, + ) + if options.loop_owner == "adapter": + adapter = implementation() + if not isinstance(adapter, RayTorchAdapter): + raise AlgorithmConfigurationError( + "adapter export requires RayTorchAdapter" + ) + runtime = TorchRuntimeContext( + algorithm_config=_source_context_config(options), + implementation_id=options.implementation_id, + world_size=descriptor.world_size, + policy_digest=descriptor.policy_digest, + execution_plan_digest=descriptor.execution_plan_digest, + run_identity=descriptor.identity, + input_bindings=options.input_bindings, + output_config=options.output_config, + input_binding_digest=options.input_binding_digest, + state_layout=descriptor.state_layout, + adapter_identity=descriptor.adapter_identity, + resume_supported=descriptor.resume_supported, + same_world_size_resume=descriptor.same_world_size_resume, + ) + from tributo.algorithms.spi import TorchStageContext + + artifact_context = TorchArtifactContext( + stage=TorchStageContext( + runtime=runtime, + stage_id=descriptor.identity.stage_id, + stage_index=options.stage_index, + is_final=True, + input_roles=options.stage_input_roles, + ), + checkpoint=checkpoint_ref, + ) + artifact_plan = adapter.artifact_plan(artifact_context) + if not isinstance(artifact_plan, TorchArtifactPlan): + raise AlgorithmConfigurationError( + "RayTorchAdapter.artifact_plan must return TorchArtifactPlan" + ) + source_context = adapter.open_export_source( + checkpoint_ref, artifact_context + ) + if not hasattr(source_context, "__enter__"): + raise AlgorithmConfigurationError( + "Adapter open_export_source must return a context manager" + ) + with source_context as source: + if not isinstance(source, ExportSource): + raise AlgorithmConfigurationError( + "Adapter export source must be an ExportSource" + ) + artifact_payload = artifact_plan.to_dict() + if source.source_kind != artifact_plan.source_kind: + raise AlgorithmConfigurationError( + "Adapter ExportSource source_kind does not match artifact plan" + ) + declared_payload = source.metadata.get("artifact_plan") + if ( + declared_payload is not None + and declared_payload != artifact_payload + ): + raise AlgorithmConfigurationError( + "Adapter ExportSource artifact plan drifted" + ) + metadata = dict(source.metadata) + metadata["artifact_plan"] = artifact_payload + checkpoint_contract = _checkpoint_contract_from_artifact_plan( + artifact_payload, + descriptor, + options, + preprocessing=source.preprocessing_state, + ) + yield source.model_copy( + update={ + "metadata": metadata, + "checkpoint_contract": checkpoint_contract, + } + ) + return + recipe = implementation() + if not isinstance(recipe, TorchRecipe): + raise AlgorithmConfigurationError("core_recipe export requires TorchRecipe") + model = _load_recipe_model(recipe, checkpoint_dir, descriptor, options) + runtime = TorchRuntimeContext( + algorithm_config=_source_context_config(options), + implementation_id=options.implementation_id, + world_size=descriptor.world_size, + policy_digest=descriptor.policy_digest, + execution_plan_digest=descriptor.execution_plan_digest, + run_identity=descriptor.identity, + input_bindings=options.input_bindings, + output_config=options.output_config, + input_binding_digest=options.input_binding_digest, + state_layout=descriptor.state_layout, + adapter_identity=descriptor.adapter_identity, + resume_supported=descriptor.resume_supported, + same_world_size_resume=descriptor.same_world_size_resume, + ) + from tributo.algorithms.spi import TorchStageContext + + artifact_context = TorchArtifactContext( + stage=TorchStageContext( + runtime=runtime, + stage_id=descriptor.identity.stage_id, + stage_index=options.stage_index, + is_final=True, + input_roles=options.stage_input_roles, + ), + checkpoint=checkpoint_ref, + ) + artifact_plan = recipe.artifact_plan(artifact_context) + if not isinstance(artifact_plan, TorchArtifactPlan): + raise AlgorithmConfigurationError( + "TorchRecipe.artifact_plan must return TorchArtifactPlan" + ) + yield _export_source( + model, + checkpoint_dir, + descriptor, + artifact_plan, + algorithm_config=options.algorithm_config, + ) + + +def _read_descriptor(root: Path) -> TorchCheckpointDescriptor: + path = root / "torch_checkpoint_descriptor.json" + if path.is_symlink() or not path.is_file(): + raise AlgorithmConfigurationError("Torch checkpoint descriptor is missing") + try: + payload = json.loads(path.read_text(encoding="utf-8")) + descriptor = TorchCheckpointDescriptor.from_dict(payload) + root_resolved = root.resolve() + commit_path = root / "torch_stage_commit.json" + if commit_path.exists() or commit_path.is_symlink(): + if commit_path.is_symlink() or not commit_path.is_file(): + raise AlgorithmConfigurationError( + "Torch checkpoint commit marker is invalid" + ) + try: + commit = json.loads(commit_path.read_text(encoding="utf-8")) + except (OSError, TypeError, ValueError) as exc: + raise AlgorithmConfigurationError( + "Torch checkpoint commit marker is malformed" + ) from exc + if not isinstance(commit, Mapping) or ( + commit.get("identity") != descriptor.identity.to_dict() + or commit.get("descriptor_digest") != descriptor.digest + ): + raise AlgorithmConfigurationError( + "Torch checkpoint commit marker does not match descriptor" + ) + actual_files: dict[str, str] = {} + for candidate in sorted(root.rglob("*")): + if candidate.is_symlink() or not candidate.resolve().is_relative_to( + root_resolved + ): + raise AlgorithmConfigurationError( + "Torch checkpoint payload escapes its root" + ) + if candidate.is_file() and candidate.name not in { + "torch_checkpoint_descriptor.json", + "torch_stage_commit.json", + ".metadata.json", + }: + actual_files[candidate.relative_to(root).as_posix()] = _sha256_file( + candidate + ) + for filename, expected_digest in descriptor.payload_files.items(): + artifact = root / filename + if artifact.is_symlink() or not artifact.resolve().is_relative_to( + root_resolved + ): + raise AlgorithmConfigurationError( + "Torch checkpoint payload escapes its root" + ) + if not artifact.is_file() or _sha256_file(artifact) != expected_digest: + raise AlgorithmConfigurationError( + "Torch checkpoint payload digest mismatch" + ) + if actual_files != dict(descriptor.payload_files): + raise AlgorithmConfigurationError( + "Torch checkpoint payload files or digest mismatch" + ) + return descriptor + except (OSError, TypeError, ValueError) as exc: + raise AlgorithmConfigurationError( + "Torch checkpoint descriptor is malformed" + ) from exc + + +def _load_recipe_model( + recipe: TorchRecipe, + root: Path, + descriptor: TorchCheckpointDescriptor, + options: TorchSourceOptions, +) -> object: + import torch + + state_path = root / "model.pt" + if state_path.is_symlink() or not state_path.is_file(): + raise AlgorithmConfigurationError("Torch export checkpoint is missing model.pt") + runtime = TorchRuntimeContext( + algorithm_config=_source_context_config(options), + implementation_id=options.implementation_id, + world_size=descriptor.world_size, + policy_digest=descriptor.policy_digest, + execution_plan_digest=descriptor.execution_plan_digest, + run_identity=descriptor.identity, + input_bindings=options.input_bindings, + output_config=options.output_config, + input_binding_digest=options.input_binding_digest, + state_layout=descriptor.state_layout, + adapter_identity=descriptor.adapter_identity, + resume_supported=descriptor.resume_supported, + same_world_size_resume=descriptor.same_world_size_resume, + ) + from tributo.algorithms.spi import ( + TorchBuildContext, + TorchModuleSet, + TorchStageContext, + ) + + modules = recipe.build_modules( + TorchBuildContext( + runtime=runtime, + stage=TorchStageContext( + runtime=runtime, + stage_id=descriptor.identity.stage_id, + stage_index=options.stage_index, + is_final=True, + input_roles=options.stage_input_roles, + ), + ) + ) + modules = ( + modules if isinstance(modules, TorchModuleSet) else TorchModuleSet(modules) + ) + model = modules["model"] + if not isinstance(model, torch.nn.Module): + raise AlgorithmConfigurationError("TorchRecipe export model is not nn.Module") + state = torch.load(state_path, map_location="cpu", weights_only=True) + if not isinstance(state, dict): + raise AlgorithmConfigurationError("Torch model checkpoint must be a state_dict") + model.load_state_dict(state) + return model + + +def _export_source( + model: object, + root: Path, + descriptor: TorchCheckpointDescriptor, + artifact_plan: object, + *, + algorithm_config: Mapping[str, Any], +) -> ExportSource: + import torch + + metrics_path = root / "metrics.json" + metrics = ( + json.loads(metrics_path.read_text(encoding="utf-8")) + if metrics_path.is_file() + else {} + ) + artifact_payload = ( + artifact_plan.to_dict() if hasattr(artifact_plan, "to_dict") else artifact_plan + ) + if not isinstance(artifact_payload, dict): + raise AlgorithmConfigurationError("Torch artifact plan is malformed") + input_signature = artifact_payload.get("input_signature", ()) + output_signature = artifact_payload.get("output_signature", ()) + if not isinstance(descriptor.identity.plan_digest, str): + raise AlgorithmConfigurationError("Torch export checkpoint has no plan digest") + checkpoint_contract = _checkpoint_contract_from_artifact_plan( + artifact_payload, + descriptor, + TorchSourceOptions( + implementation_ref="internal:torch", + implementation_code_digest=descriptor.implementation_code_digest, + implementation_id=descriptor.identity.implementation_id, + policy_digest=descriptor.policy_digest, + plan_digest=descriptor.identity.plan_digest, + input_binding_digest=descriptor.input_binding_digest, + algorithm_config=dict(algorithm_config), + ), + ) + sample_inputs: dict[str, object] = {} + for field in input_signature: + if not isinstance(field, dict): + raise AlgorithmConfigurationError("Torch input signature is malformed") + shape = tuple( + 1 if dim == "batch" else int(dim) for dim in field.get("shape", (1,)) + ) + dtype = getattr(torch, str(field.get("dtype", "float32")), None) + if dtype is None: + raise AlgorithmConfigurationError( + "Torch artifact input dtype is unsupported" + ) + field_name = str(field["name"]) + sample_inputs[field_name] = ( + torch.ones(shape, dtype=dtype) + if ( + field_name == "input_ids" + and dtype + in { + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + } + ) + else torch.zeros(shape, dtype=dtype) + ) + return ExportSource( + source_kind="torch_module", + model_object=model, + architecture_id=descriptor.identity.implementation_id, + model_config_data=dict(algorithm_config.get("model", {})) + if isinstance(algorithm_config.get("model", {}), Mapping) + else {}, + feature_schema={ + "input_signature": input_signature, + "output_signature": output_signature, + "input_names": [ + field["name"] + for field in input_signature + if isinstance(field, Mapping) and isinstance(field.get("name"), str) + ], + }, + preprocessing_state={}, + sample_inputs=sample_inputs, + checkpoint_contract=checkpoint_contract, + metadata={ + "framework": "pytorch", + "torch_runtime_api_version": 1, + "artifact_plan": artifact_payload, + "metrics": metrics, + }, + source_fingerprint=_sha256_file(root / "model.pt"), + ) + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +__all__ = ["RayTorchSourceProvider", "TorchSourceOptions"] diff --git a/src/tributo/integrations/sources/ray_torch_recipe.py b/src/tributo/integrations/sources/ray_torch_recipe.py deleted file mode 100644 index 6225f0d..0000000 --- a/src/tributo/integrations/sources/ray_torch_recipe.py +++ /dev/null @@ -1,194 +0,0 @@ -"""Ray Torch recipe checkpoint to the existing ExportSource contract.""" - -from __future__ import annotations - -import hashlib -import json -from contextlib import contextmanager -from pathlib import Path -from typing import Any, ClassVar, Generator - -from pydantic import BaseModel, ConfigDict, Field - -from tributo.algorithms.api import QualifiedReference -from tributo.algorithms.core.worker import _load_reference, _validate_module_digest -from tributo.algorithms.spi import TorchTrainingRecipe, TrainingRecipeV2 -from tributo.exporting.models import ExportCheckpointV1, ExportSource -from tributo.training.checkpoint import checkpoint_directory -from tributo.util.annotations import PublicAPI - - -@PublicAPI(stability="alpha") -class TorchRecipeSourceOptions(BaseModel): - """Bind a checkpoint to the exact trusted recipe implementation.""" - - model_config = ConfigDict(frozen=True, extra="forbid") - - recipe_ref: str = Field(min_length=1) - recipe_code_digest: str | None = Field( - default=None, - pattern=r"^[0-9a-f]{64}$", - ) - implementation_id: str = Field(min_length=1) - - -@PublicAPI(stability="alpha") -class RayTorchRecipeSourceProvider: - """Reconstruct a recipe model for the existing Torch ONNX exporter.""" - - api_version: ClassVar[int] = 1 - provider_id: ClassVar[str] = "ray-torch-recipe-v1" - trainer_type: ClassVar[str] = "torch_recipe" - priority: ClassVar[int] = 100 - - def open_source( - self, - result: Any, - config: BaseModel | None = None, - ) -> Any: - """Open one verified Ray Train recipe checkpoint.""" - options = TorchRecipeSourceOptions.model_validate( - config.model_dump() if config is not None else {} - ) - return _open_recipe_source(result, options) - - -@contextmanager -def _open_recipe_source( - result: Any, - options: TorchRecipeSourceOptions, -) -> Generator[ExportSource, None, None]: - checkpoint = getattr(result, "checkpoint", result) - if checkpoint is None: - raise ValueError("Torch recipe training result has no checkpoint") - with checkpoint_directory(checkpoint) as checkpoint_dir: - yield _build_source(checkpoint_dir, options) - - -def _build_source( - checkpoint_dir: Path, - options: TorchRecipeSourceOptions, -) -> ExportSource: - import torch - - root = checkpoint_dir.resolve() - model_path = checkpoint_dir / "model.pt" - config_path = checkpoint_dir / "model_config.json" - for path in (model_path, config_path): - if path.is_symlink() or not path.resolve().is_relative_to(root): - raise ValueError("Torch recipe checkpoint artifact escapes its root") - if not path.is_file(): - raise FileNotFoundError(f"Torch recipe checkpoint is missing {path.name!r}") - payload = json.loads(config_path.read_text(encoding="utf-8")) - contract = ExportCheckpointV1.model_validate(payload) - if contract.trainer_type != "torch_recipe": - raise ValueError("checkpoint is not a Torch recipe export source") - if payload.get("recipe_ref") != options.recipe_ref: - raise ValueError("checkpoint recipe reference does not match the plan") - if payload.get("recipe_code_digest") != options.recipe_code_digest: - raise ValueError("checkpoint recipe code digest does not match the plan") - if contract.architecture_id != options.implementation_id: - raise ValueError("checkpoint implementation identity does not match the plan") - - reference = QualifiedReference.parse(options.recipe_ref) - _validate_module_digest(reference, options.recipe_code_digest) - recipe_cls = _load_reference(reference) - if not isinstance(recipe_cls, type) or not issubclass( - recipe_cls, (TorchTrainingRecipe, TrainingRecipeV2) - ): - raise ValueError( - "recipe reference does not resolve to TorchTrainingRecipe or " - "TrainingRecipeV2" - ) - try: - recipe = recipe_cls() - except TypeError as exc: - raise ValueError("Torch recipe must have a no-argument constructor") from exc - model_config = payload.get("model", {}) - if not isinstance(model_config, dict): - raise ValueError("Torch recipe model config must be a JSON object") - if isinstance(recipe, TrainingRecipeV2): - modules = recipe.build_modules({"model": model_config}) - if not isinstance(modules, dict) or "model" not in modules: - raise ValueError("TrainingRecipeV2 build_modules did not provide model") - model = modules["model"] - else: - model = recipe.model_factory(model_config) - if not isinstance(model, torch.nn.Module): - raise ValueError("Torch recipe model_factory did not return nn.Module") - state = torch.load(model_path, map_location="cpu", weights_only=True) - if not isinstance(state, dict): - raise ValueError("Torch recipe model.pt must contain one state_dict") - model.load_state_dict(state) - input_names = [field.name for field in contract.input_schema] - model = _wrap_dense_columns(model, input_names) - sample_inputs = { - field.name: torch.zeros( - tuple(2 if value == "batch" else int(value) for value in field.shape), - dtype=torch.float32, - ) - for field in contract.input_schema - } - metrics_path = checkpoint_dir / "metrics.json" - metrics = ( - json.loads(metrics_path.read_text(encoding="utf-8")) - if metrics_path.is_file() - else {} - ) - return ExportSource( - source_kind="torch_module", - model_object=model, - architecture_id=contract.architecture_id, - model_config_data=model_config, - feature_schema={ - "feature_names": input_names, - "input_schema": [ - field.model_dump(mode="json") for field in contract.input_schema - ], - }, - preprocessing_state={}, - sample_inputs=sample_inputs, - metadata={ - "framework": contract.framework, - "framework_version": contract.framework_version, - "task_type": contract.task_type, - "trainer_type": contract.trainer_type, - "metrics": metrics, - }, - source_fingerprint=_sha256_file(model_path)[:16], - checkpoint_contract=contract, - ) - - -def _sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as stream: - while chunk := stream.read(1024 * 1024): - digest.update(chunk) - return digest.hexdigest() - - -def _wrap_dense_columns(model: Any, input_names: list[str]) -> Any: - import torch - - class _DenseColumnModule(torch.nn.Module): - def __init__(self, wrapped: torch.nn.Module) -> None: - super().__init__() - self.wrapped = wrapped - - def forward(self, *columns: Any) -> Any: - if not columns: - raise ValueError("Torch recipe export requires input columns") - rows = columns[0].shape[0] - features = torch.cat( - [column.reshape(rows, -1).float() for column in columns], - dim=1, - ) - return self.wrapped(features) - - if not isinstance(model, torch.nn.Module) or not input_names: - raise ValueError("Torch recipe export requires a model and named inputs") - return _DenseColumnModule(model) - - -__all__ = ["RayTorchRecipeSourceProvider", "TorchRecipeSourceOptions"] diff --git a/src/tributo/plugin.py b/src/tributo/plugin.py index 9b18faf..01739df 100644 --- a/src/tributo/plugin.py +++ b/src/tributo/plugin.py @@ -242,6 +242,7 @@ def validate_distributed_algorithm_descriptor( DistributedAlgorithmDescriptor, DistributionStrategy, InputDistribution, + TorchPolicy, ) from tributo.algorithms.api.models import FORMAL_DISTRIBUTED_STRATEGY_CONTRACTS from tributo.algorithms.spi import ( @@ -251,8 +252,8 @@ def validate_distributed_algorithm_descriptor( JoblibEstimatorRecipe, MapReduceAlgorithm, ParallelEnsembleAlgorithm, - TorchTrainingRecipe, - TrainingRecipeV2, + RayTorchAdapter, + TorchRecipe, ) if not isinstance(descriptor, DistributedAlgorithmDescriptor): @@ -294,15 +295,13 @@ def validate_distributed_algorithm_descriptor( DistributionStrategy.RAY_ITERATIVE_OPTIMIZATION: ( IterativeOptimizationAlgorithm ), - DistributionStrategy.RAY_TRAIN_RECIPE_V2: TrainingRecipeV2, + DistributionStrategy.RAY_TRAIN_TORCH: TorchRecipe, }[distribution_spec.strategy] - if distribution_spec.strategy is DistributionStrategy.RAY_TRAIN_COLLECTIVE and str( - implementation_descriptor.executable_factory_ref - ) == ( - "tributo.integrations.algorithm_runtimes.torch_recipe:" - "create_torch_recipe_algorithm" - ): - expected_base = TorchTrainingRecipe + if distribution_spec.strategy is DistributionStrategy.RAY_TRAIN_TORCH: + torch_policy = cast(TorchPolicy, distribution_spec.policy) + expected_base = ( + RayTorchAdapter if torch_policy.loop_owner == "adapter" else TorchRecipe + ) contract = FORMAL_DISTRIBUTED_STRATEGY_CONTRACTS[distribution_spec.strategy] if implementation_descriptor.execution_mode is not contract.execution_mode: raise ValueError("implementation execution mode conflicts with strategy") diff --git a/tests/algorithms/test_official_algorithm_migration.py b/tests/algorithms/test_official_algorithm_migration.py index cb0c863..18b0a57 100644 --- a/tests/algorithms/test_official_algorithm_migration.py +++ b/tests/algorithms/test_official_algorithm_migration.py @@ -35,7 +35,7 @@ def test_core_composition_root_contains_only_algorithm_neutral_plugins() -> None } assert {item.validator_id for item in validators} == {"onnx-runtime-v1"} assert {item.provider_id for item in first_party_source_providers()} == { - "ray-torch-recipe-v1" + "ray-torch-v1" } assert {item.flavor_id for item in first_party_model_flavors()} == { "onnx-runtime-v1" diff --git a/tests/algorithms/test_portable_models.py b/tests/algorithms/test_portable_models.py index f520a1a..c10522a 100644 --- a/tests/algorithms/test_portable_models.py +++ b/tests/algorithms/test_portable_models.py @@ -64,7 +64,6 @@ def _formal_policy( ): if strategy in { DistributionStrategy.RAY_TRAIN_COLLECTIVE, - DistributionStrategy.RAY_TRAIN_RECIPE_V2, }: return CollectivePolicy( backend="gloo", @@ -109,7 +108,6 @@ def _formal_descriptor( DistributionStrategy.RAY_ITERATIVE_OPTIMIZATION: ( ExecutionMode.ITERATIVE_OPTIMIZATION ), - DistributionStrategy.RAY_TRAIN_RECIPE_V2: ExecutionMode.TRAINING_RECIPE_V2, }[strategy] return AlgorithmBuilder.from_distributed_algorithm( spec=make_spec( @@ -383,6 +381,8 @@ def test_distributed_builder_lowers_each_strategy_deterministically( def test_distributed_builder_contract_is_shared_with_registration_fields() -> None: for strategy, contract in FORMAL_DISTRIBUTED_STRATEGY_CONTRACTS.items(): + if strategy is DistributionStrategy.RAY_TRAIN_TORCH: + continue registration = _formal_descriptor(strategy).registration distribution = registration.distribution_spec assert distribution is not None diff --git a/tests/algorithms/test_torch_recipe_contract.py b/tests/algorithms/test_torch_recipe_contract.py index 8562e4c..94d889e 100644 --- a/tests/algorithms/test_torch_recipe_contract.py +++ b/tests/algorithms/test_torch_recipe_contract.py @@ -4,7 +4,7 @@ import pytest -from tributo.algorithms import AlgorithmBuilder, TorchTrainingRecipe +from tributo.algorithms import AlgorithmBuilder, TorchRecipe from tributo.algorithms.api import ( AlgorithmConfigurationError, EnvironmentSpec, @@ -20,11 +20,11 @@ def _descriptor(): - return AlgorithmBuilder.from_torch_recipe( + return AlgorithmBuilder.from_torch( spec=make_spec( "external_torch_recipe", operations=("fit",), - mode=ExecutionMode.COLLECTIVE, + mode=ExecutionMode.RAY_TRAIN_TORCH, default_config={ "model": {"input_features": 2}, "output": {"bundle_uri": "./bundle"}, @@ -51,37 +51,40 @@ def _descriptor(): package_name="example-torch-recipe", package_version="1.0.0", tributo_version_spec=">=1,<2", + code_digest="0" * 64, + descriptor_api_version=1, ) -def test_torch_recipe_is_narrower_than_collective_worker_loop() -> None: - abstract_methods = TorchTrainingRecipe.__abstractmethods__ +def test_torch_recipe_has_only_typed_math_and_export_hooks() -> None: + abstract_methods = TorchRecipe.__abstractmethods__ assert abstract_methods == { - "loss_factory", - "metric_factories", - "model_factory", - "optimizer_factory", + "adapt_batch", + "artifact_plan", + "build_modules", + "configure_optimizers", + "metric_plan", + "training_step", + "validation_step", } - assert "train_loop_per_worker" not in abstract_methods - assert "checkpoint_state" not in abstract_methods + assert "execution_plan" not in abstract_methods -def test_builder_lowers_recipe_to_existing_collective_runtime() -> None: +def test_builder_lowers_recipe_to_unified_torch_runtime() -> None: descriptor = _descriptor() registration = descriptor.registration - assert registration.implementation.runtime_id == "tributo.ray_train_collective" + assert registration.implementation.runtime_id == "tributo.ray_train_torch" assert str(registration.implementation.implementation_ref) == ( "tests.support.torch_recipe:BinaryLinearRecipe" ) assert str(registration.implementation.executable_factory_ref) == ( - "tributo.integrations.algorithm_runtimes.torch_recipe:" - "create_torch_recipe_algorithm" + "tributo.integrations.algorithm_runtimes.ray_train_torch:create_torch_algorithm" ) assert str(registration.implementation.exporter_ref) == ( - "tributo.integrations.algorithm_runtimes.torch_recipe:" - "export_torch_recipe_result" + "tributo.integrations.algorithm_runtimes.ray_train_torch:" + "export_ray_train_torch_result" ) assert registration.distribution_spec is not None assert registration.distribution_spec.result_policy is ResultPolicy.BUNDLE_REQUIRED @@ -90,6 +93,7 @@ def test_builder_lowers_recipe_to_existing_collective_runtime() -> None: "train_loss": MetricReduction.SUM_COUNT, } assert registration.implementation.allowed_config_keys == ( + "data", "loss", "metrics", "model", @@ -102,11 +106,11 @@ def test_builder_lowers_recipe_to_existing_collective_runtime() -> None: def test_builder_rejects_train_loss_reducer_override() -> None: with pytest.raises(AlgorithmConfigurationError, match="train_loss"): - AlgorithmBuilder.from_torch_recipe( + AlgorithmBuilder.from_torch( spec=make_spec( "invalid_torch_recipe", operations=("fit",), - mode=ExecutionMode.COLLECTIVE, + mode=ExecutionMode.RAY_TRAIN_TORCH, ), implementation_id="example.torch_recipe.invalid", implementation_version="1.0.0", @@ -122,4 +126,5 @@ def test_builder_rejects_train_loss_reducer_override() -> None: package_name="example-torch-recipe", package_version="1.0.0", tributo_version_spec=">=1,<2", + code_digest="0" * 64, ) diff --git a/tests/algorithms/test_torch_recipe_plugin.py b/tests/algorithms/test_torch_recipe_plugin.py index 922d119..a8d026c 100644 --- a/tests/algorithms/test_torch_recipe_plugin.py +++ b/tests/algorithms/test_torch_recipe_plugin.py @@ -6,7 +6,7 @@ import sys from pathlib import Path -from tributo.algorithms import TorchTrainingRecipe +from tributo.algorithms import TorchRecipe from tributo.algorithms.api import ExecutionProfile, ResultPolicy @@ -37,7 +37,7 @@ def test_out_of_tree_recipe_descriptor_uses_public_lowering_contract() -> None: finally: sys.path.remove(str(_fixture_source())) - assert issubclass(module.ThirdPartyBinaryLinearRecipe, TorchTrainingRecipe) + assert issubclass(module.ThirdPartyBinaryLinearRecipe, TorchRecipe) distribution = module.DESCRIPTOR.registration.distribution_spec assert distribution is not None assert distribution.result_policy is ResultPolicy.BUNDLE_REQUIRED diff --git a/tests/algorithms/test_torch_recipe_worker.py b/tests/algorithms/test_torch_recipe_worker.py index f37d248..bce4ff8 100644 --- a/tests/algorithms/test_torch_recipe_worker.py +++ b/tests/algorithms/test_torch_recipe_worker.py @@ -10,14 +10,75 @@ import ray import ray.train -from tests.support.torch_recipe import BinaryLinearRecipe -from tributo.integrations.algorithm_runtimes.torch_recipe import ( +from tributo.algorithms import ( + TorchBatch, + TorchLossContribution, + TorchMetricPlan, + TorchModuleSet, + TorchOptimizationPlan, + TorchRecipe, + TorchRuntimeContext, + TorchStageContext, + TorchStageRunIdentity, + TorchStepResult, +) +from tributo.integrations.algorithm_runtimes.ray_train_torch import ( torch_recipe_train_loop_per_worker, ) torch = pytest.importorskip("torch") +class BinaryLinearRecipe(TorchRecipe): + def build_modules(self, context: object) -> TorchModuleSet: + return TorchModuleSet( + {"model": torch.nn.Linear(2, 1), "loss": torch.nn.MSELoss()} + ) + + def adapt_batch(self, batch: object, context: object) -> TorchBatch: + del context + features = torch.stack((batch["x1"], batch["x2"]), dim=1) + return TorchBatch( + positional=(features,), + targets=batch["label"].reshape(-1, 1), + local_rows=len(batch["label"]), + ) + + def training_step( + self, modules: TorchModuleSet, batch: TorchBatch, context: object + ) -> TorchStepResult: + del context + predictions = modules["model"](batch.positional[0]) + numerator = torch.nn.functional.mse_loss( + predictions, batch.targets, reduction="sum" + ) + return TorchStepResult( + outputs={"prediction": predictions}, + loss=TorchLossContribution(numerator, batch.local_rows), + ) + + def validation_step( + self, modules: TorchModuleSet, batch: TorchBatch, context: object + ) -> TorchStepResult: + return self.training_step(modules, batch, context) + + def configure_optimizers( + self, modules: TorchModuleSet, context: object + ) -> TorchOptimizationPlan: + del context + return TorchOptimizationPlan( + torch.optim.SGD(modules["model"].parameters(), lr=0.1) + ) + + def metric_plan(self, context: TorchRuntimeContext) -> TorchMetricPlan: + del context + return TorchMetricPlan({"train_loss": "sum_count"}) + + def artifact_plan(self, context: object) -> dict[str, object]: + del context + return {"source_kind": "torch_module"} + + class _Iterator: def __init__(self, batches: list[dict[str, Any]]) -> None: self._batches = batches @@ -47,159 +108,127 @@ def get_assigned_resources(self) -> dict[str, float]: return {"CPU": 1.0} -def test_default_recipe_streams_torch_batches_and_reports_resumable_checkpoint( +def test_recipe_worker_reports_typed_checkpoint_and_exact_coverage( monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: - batches = [ - { - "x1": torch.tensor([0.0, 1.0, 0.0]), - "x2": torch.tensor([0.0, 0.0, 1.0]), - "label": torch.tensor([0.0, 1.0, 1.0]), - }, - { - "x1": torch.tensor([1.0, 1.0]), - "x2": torch.tensor([0.0, 1.0]), - "label": torch.tensor([1.0, 1.0]), - }, - ] - data = _Iterator(batches) - validation = _Iterator( + data = _Iterator( [ { "x1": torch.tensor([0.0, 1.0]), "x2": torch.tensor([0.0, 1.0]), "label": torch.tensor([0.0, 1.0]), - } + }, + { + "x1": torch.tensor([1.0]), + "x2": torch.tensor([0.0]), + "label": torch.tensor([1.0]), + }, ] ) reports: list[tuple[dict[str, Any], set[str]]] = [] def report(metrics: dict[str, Any], checkpoint: Any | None = None) -> None: - files: set[str] = set() - if checkpoint is not None: - with checkpoint.as_directory() as raw_directory: - files = { - path.name - for path in Path(raw_directory).iterdir() - if path.is_file() - } - reports.append((dict(metrics), files)) - - def get_dataset_shard(name: str) -> _Iterator: - if name == "train": - return data - if name == "val": - return validation - raise KeyError(name) + assert checkpoint is not None + with checkpoint.as_directory() as directory: + reports.append( + (dict(metrics), {path.name for path in Path(directory).iterdir()}) + ) + monkeypatch.setattr(ray.train.torch, "prepare_model", lambda model: model) monkeypatch.setattr(ray.train, "get_context", lambda: _TrainContext()) - monkeypatch.setattr(ray.train, "get_dataset_shard", get_dataset_shard) + monkeypatch.setattr( + ray.train, + "get_dataset_shard", + lambda name: data if name == "train" else (_ for _ in ()).throw(KeyError(name)), + ) monkeypatch.setattr(ray.train, "get_checkpoint", lambda: None) monkeypatch.setattr(ray.train, "report", report) monkeypatch.setattr(ray, "get_runtime_context", lambda: _RuntimeContext()) - - torch_recipe_train_loop_per_worker( - { - "model": {"input_features": 2}, - "optimizer": {"learning_rate": 0.1}, - "training": { - "epochs": 1, - "batch_size": 3, - "prefetch_batches": 2, - "seed": 7, - }, - "ray": {"resume": {"checkpoint_interval": 1}}, - "_tributo_recipe_ref": ("tests.support.torch_recipe:BinaryLinearRecipe"), - "_tributo_recipe_code_digest": None, - "_tributo_implementation_id": "example.binary_linear", - "_tributo_algorithm": "binary_linear", - "_tributo_feature_names": ["x1", "x2"], - "_tributo_label_name": "label", - "_tributo_weight_name": None, - "_tributo_input_binding_digest": "a" * 64, - "_tributo_distribution_spec_digest": "b" * 64, - "_tributo_metric_reducers": { - "accuracy": "sum_count", - "train_loss": "sum_count", - }, - }, - BinaryLinearRecipe(), + identity = TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "example", + "example.binary", + "0" * 64, + "1" * 64, + "2" * 64, ) - - assert data.calls == [ + runtime = TorchRuntimeContext( + {}, + "example.binary", + 1, + "1" * 64, + "2" * 64, + identity, + input_binding_digest="3" * 64, + ) + stage = TorchStageContext(runtime, "train", 0, True, ("train",)) + torch_recipe_train_loop_per_worker( { - "batch_size": 3, - "prefetch_batches": 2, - "dtypes": torch.float32, - "drop_last": False, - "local_shuffle_buffer_size": None, - "local_shuffle_seed": 7, + "training": {"epochs": 1, "batch_size": 2}, + "_core_implementation_ref": "tests.algorithms.test_torch_recipe_worker:BinaryLinearRecipe", + "_core_implementation_code_digest": "0" * 64, + "_core_input_binding_digest": "3" * 64, + "_core_stage_context": stage.to_dict(), } - ] - assert len(reports) == 1 - metrics, files = reports[0] - assert metrics["epoch"] == 1 - assert 0.0 <= metrics["val_accuracy"] <= 1.0 - assert metrics["val_loss"] >= 0.0 - assert metrics["metric_reducers"] == { - "accuracy": "sum_count", - "train_loss": "sum_count", - } - assert metrics["execution_workers"][0]["input_rows"] == { - "train": 5, - "val": 2, - } - assert metrics["execution_workers"][0]["batch_count"] == 2 - assert files == { - "metrics.json", + ) + assert reports and reports[0][0]["checkpoint_descriptor"]["completed_step"] == 2 + assert { "model.pt", - "model_config.json", "optimizer.pt", - "resume.json", - "rng_state.json", "scaler.pt", - "training_state.json", - } - assert validation.calls == [ - { - "batch_size": 3, - "prefetch_batches": 2, - "dtypes": torch.float32, - "drop_last": False, - } - ] + "rng_state.pt", + "torch_checkpoint_descriptor.json", + } <= reports[0][1] -def test_default_recipe_resumes_from_explicit_worker_visible_checkpoint( +def test_recipe_worker_rejects_stale_retry_checkpoint( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - from tributo.training.checkpoint import ( - capture_rng_state, - write_resume_manifest, - ) + from tributo.algorithms import TorchCheckpointDescriptor recipe = BinaryLinearRecipe() - model = recipe.model_factory({"input_features": 2}) - assert isinstance(model, torch.nn.Module) - optimizer = recipe.optimizer_factory(model, {"learning_rate": 0.1}) + modules = recipe.build_modules(None) + model = modules["model"] + optimizer = recipe.configure_optimizers(modules, None).optimizer torch.save(model.state_dict(), tmp_path / "model.pt") torch.save(optimizer.state_dict(), tmp_path / "optimizer.pt") - (tmp_path / "rng_state.json").write_text( - json.dumps({"rank_states": [capture_rng_state()]}), - encoding="utf-8", + torch.save({}, tmp_path / "scaler.pt") + (tmp_path / "rng_state.pt").write_bytes(torch.get_rng_state().numpy().tobytes()) + identity = TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "example", + "example.binary", + "0" * 64, + "1" * 64, + "2" * 64, ) - write_resume_manifest( - tmp_path, - trainer_type="torch_recipe", - completed_step=1, - framework="pytorch", - framework_version=torch.__version__, - payload_files=("model.pt", "optimizer.pt", "rng_state.json"), - payload_metadata={ - "world_size": 1, - "distribution_spec_digest": "b" * 64, - }, + payload_files = { + name: __import__("hashlib").sha256((tmp_path / name).read_bytes()).hexdigest() + for name in ("model.pt", "optimizer.pt", "scaler.pt", "rng_state.pt") + } + descriptor = TorchCheckpointDescriptor( + 1, + identity, + identity.run_config_name, + "replicated", + 1, + 1, + "1" * 64, + "2" * 64, + "3" * 64, + "0" * 64, + payload_files, + ) + (tmp_path / "torch_checkpoint_descriptor.json").write_text( + json.dumps(descriptor.to_dict()), encoding="utf-8" ) data = _Iterator( [ @@ -219,7 +248,19 @@ def get_dataset_shard(name: str) -> _Iterator: monkeypatch.setattr(ray.train, "get_context", lambda: _TrainContext()) monkeypatch.setattr(ray.train, "get_dataset_shard", get_dataset_shard) - monkeypatch.setattr(ray.train, "get_checkpoint", lambda: None) + monkeypatch.setattr(ray.train.torch, "prepare_model", lambda model: model) + + class _Checkpoint: + def as_directory(self): + from contextlib import contextmanager + + @contextmanager + def opened(): + yield tmp_path + + return opened() + + monkeypatch.setattr(ray.train, "get_checkpoint", lambda: _Checkpoint()) monkeypatch.setattr( ray.train, "report", @@ -247,8 +288,35 @@ def get_dataset_shard(name: str) -> _Iterator: "accuracy": "sum_count", "train_loss": "sum_count", }, + "_core_implementation_ref": "tests.support.torch_recipe:BinaryLinearRecipe", + "_core_implementation_code_digest": "0" * 64, + "_core_input_binding_digest": "a" * 64, + "_core_stage_context": TorchStageContext( + TorchRuntimeContext( + {}, + "example.binary_linear", + 1, + "1" * 64, + "2" * 64, + TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "binary", + "example.binary_linear", + "0" * 64, + "1" * 64, + "2" * 64, + ), + input_binding_digest="a" * 64, + ), + "train", + 0, + True, + ("train",), + ).to_dict(), }, - recipe, ) - assert [report["epoch"] for report in reports] == [2] + assert reports and reports[0]["checkpoint_descriptor"]["completed_step"] == 3 diff --git a/tests/algorithms/test_torch_runtime_contract_v1.py b/tests/algorithms/test_torch_runtime_contract_v1.py new file mode 100644 index 0000000..da3041a --- /dev/null +++ b/tests/algorithms/test_torch_runtime_contract_v1.py @@ -0,0 +1,920 @@ +"""Focused Core tests for the unified Torch v1 public contract.""" + +from __future__ import annotations + +import inspect +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest + +from tributo.algorithms.api import ( + AlgorithmConfigurationError, + AlgorithmExecutionError, + DistributionStrategy, + InputBinding, + MetricReduction, + ResultPolicy, + SingleStageTorchPlan, + TorchAccumulationWindow, + TorchBackwardContext, + TorchCheckpointDescriptor, + TorchCheckpointLocator, + TorchCheckpointProgress, + TorchCompositeLossContribution, + TorchDatasetRoute, + TorchGlobalLossReduction, + TorchLossContribution, + TorchMetricContribution, + TorchMetricPolicy, + TorchMetricReductionContext, + TorchPolicy, + TorchPreflightLease, + TorchPreflightTokenData, + TorchRankProgressStatistics, + TorchRecoveryEnvelope, + TorchStageRunIdentity, + TorchStageSpec, + apply_torch_loss_backward, + reduce_torch_metrics, + report_torch_checkpoint, + torch_run_config_name, +) +from tributo.algorithms.spi import TorchRuntimeContext, TorchStageContext + + +class _Scalar: + ndim = 0 + + def __init__(self, value: float) -> None: + self.value = value + + def detach(self) -> "_Scalar": + return self + + def item(self) -> float: + return self.value + + +def _identity_kwargs() -> dict[str, object]: + return { + "run_id": "aabbccdd", + "invocation_id": "11223344", + "algorithm": "example", + "implementation_ref": "example:Recipe", + "implementation_code_digest": "0" * 64, + "policy_digest": "1" * 64, + "execution_plan_digest": "2" * 64, + "runtime_id": "tributo.ray_train_torch", + "plan_digest": "3" * 64, + } + + +def test_torch_policy_and_run_name_are_deterministic() -> None: + route = TorchDatasetRoute("train", "split_exact") + execution_plan = SingleStageTorchPlan( + stage=TorchStageSpec("train", "example:loop", ("train",)) + ) + policy = TorchPolicy( + torch_runtime_api_version=1, + loop_owner="core_recipe", + parallelism_id="torch.ddp.replicated", + dataset_routing=(route,), + execution_plan=execution_plan, + state_layout="replicated", + metric_reducers={"train_loss": MetricReduction.SUM_COUNT}, + ) + assert policy.digest == TorchPolicy.from_dict(policy.to_dict()).digest + identity = TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "example", + "example.recipe", + "0" * 64, + policy.digest, + execution_plan.digest, + ) + assert torch_run_config_name(identity) == identity.run_config_name + + +def test_preflight_lease_is_one_shot_and_identity_bound() -> None: + lease = TorchPreflightLease(TorchPreflightTokenData(**_identity_kwargs())) + lease.claim( + run_id="aabbccdd", + invocation_id="11223344", + plan_digest="3" * 64, + runtime_id="tributo.ray_train_torch", + ) + lease.consume( + run_id="aabbccdd", + invocation_id="11223344", + plan_digest="3" * 64, + runtime_id="tributo.ray_train_torch", + ) + with pytest.raises(AlgorithmExecutionError): + lease.consume( + run_id="aabbccdd", + invocation_id="11223344", + plan_digest="3" * 64, + runtime_id="tributo.ray_train_torch", + ) + + +def test_loss_contribution_requires_zero_dimensional_scalar() -> None: + assert TorchLossContribution(_Scalar(2.0), 3).normalizer == 3.0 + with pytest.raises(AlgorithmConfigurationError): + TorchLossContribution(2.0, 3) + composite = TorchCompositeLossContribution( + "schema", + {"loss_a": _Scalar(2.0), "loss_b": _Scalar(1.0), "loss_c": _Scalar(3.0)}, + {"count_a": 2, "count_b": 4}, + ) + assert set(composite.differentiable_components) == {"loss_a", "loss_b", "loss_c"} + + +def test_backward_and_metric_helpers_use_explicit_normalizers() -> None: + events: list[object] = [] + result = apply_torch_loss_backward( + TorchLossContribution(_Scalar(1.0), 2), + TorchAccumulationWindow(index=0, expected_micro_batches=1), + TorchBackwardContext( + world_size=2, + backward=lambda value: events.append(("backward", value)), + reduce_normalizer=lambda value: value * 3, + finalize_window=lambda scale: events.append(("scale", scale)), + ), + ) + assert result.global_normalizer == 6 + assert events[-1] == ("scale", 2 / 6) + metrics = reduce_torch_metrics( + {"loss": TorchMetricContribution(4, 2)}, + TorchMetricPolicy({"loss": "sum_count"}), + TorchMetricReductionContext( + lambda name, value, reducer: value.numerator / value.normalizer + ), + ) + assert metrics.values == {"loss": 2.0} + + +def test_backward_helper_reduces_only_the_accumulation_window_total() -> None: + reductions: list[float] = [] + scales: list[float] = [] + context = TorchBackwardContext( + world_size=2, + backward=lambda value: None, + reduce_normalizer=lambda value: reductions.append(value) or value * 2, + finalize_window=scales.append, + ) + first = apply_torch_loss_backward( + TorchLossContribution(_Scalar(1), 2), + TorchAccumulationWindow(0, 2), + context, + ) + assert not first.window_complete + assert reductions == [] + second = apply_torch_loss_backward( + TorchLossContribution(_Scalar(1), 3), + TorchAccumulationWindow(0, 2, 1, first.global_normalizer), + context, + ) + assert second.window_complete + assert reductions == [5] + assert scales == [2 / 10] + + +def test_locator_rejects_local_paths_and_policy_replicate_budget_is_explicit() -> None: + with pytest.raises(AlgorithmConfigurationError): + TorchCheckpointLocator("/tmp/checkpoint", "0" * 64) + route = TorchDatasetRoute( + "nodes", "replicate", max_rows=10, max_bytes_per_worker=10 + ) + plan = SingleStageTorchPlan( + stage=TorchStageSpec("train", "example:loop", ("nodes",)) + ) + with pytest.raises(AlgorithmConfigurationError): + TorchPolicy( + 1, + "core_recipe", + "torch.ddp.replicated", + (route,), + plan, + "replicated", + {"train_loss": MetricReduction.SUM_COUNT}, + ) + + +def test_composite_global_state_keeps_component_and_normalizer_names_independent() -> ( + None +): + from tributo.algorithms.api import TorchCompositeGlobalState + + state = TorchCompositeGlobalState( + components={"positive": 2.0, "negative": -1.0}, + normalizers={"positive_count": 3.0, "negative_count": 4.0}, + ) + assert set(state.components) == {"positive", "negative"} + assert set(state.normalizers) == {"positive_count", "negative_count"} + + +def test_stage_dependency_is_allowed_when_external_recovery_is_disabled() -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _control_for_stage, + ) + + control = _control_for_stage( + SimpleNamespace( + algorithm_config={}, + runtime=SimpleNamespace(resume_from=None), + ), + SimpleNamespace( + resume_supported=False, + digest="1" * 64, + execution_plan=SimpleNamespace(digest="2" * 64), + ), + SimpleNamespace(stage_id="student", checkpoint_from_stage="teacher"), + run_id="aabbccdd", + invocation_id="11223344", + predecessor={ + "locator": "s3://bucket/teacher-checkpoint", + "descriptor_digest": "3" * 64, + }, + ) + assert control is not None + assert control["purpose"] == "stage_dependency" + assert control["source_stage_id"] == "teacher" + + +def test_role_evidence_falls_back_to_primary_binding_for_alias_roles() -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _binding_digest_for_role, + ) + + class Descriptors: + def get(self, role: str) -> object: + raise AlgorithmConfigurationError(f"unknown resolved input role: {role}") + + primary = SimpleNamespace(binding_digest="3" * 64) + plan = SimpleNamespace( + input_descriptors=Descriptors(), primary_input_descriptor=primary + ) + assert _binding_digest_for_role(plan, "val") == "3" * 64 + + +def test_worker_evidence_defaults_only_missing_declared_resources() -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _normalize_worker_evidence, + ) + + plan = SimpleNamespace( + runtime=SimpleNamespace( + num_cpus=2.0, + num_gpus=1.0, + custom_resources={"accelerator": 1.0}, + memory_bytes=1024, + ) + ) + existing = {"num_cpus": 9.0, "num_gpus": 0.0, "custom": {}} + records = _normalize_worker_evidence( + [{"worker_id": "a"}, {"worker_id": "b", "resources": existing}], + plan, + ) + assert records[0]["resources"] == { + "num_cpus": 2.0, + "num_gpus": 1.0, + "custom": {"accelerator": 1.0}, + "memory_bytes": 1024, + } + assert records[1]["resources"] is existing + + +def test_component_state_details_project_stage_coverage() -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _component_state_details, + ) + + stages = ( + SimpleNamespace( + stage_id="pretrain", + state_digest="a" * 64, + roles=(SimpleNamespace(role="train", present=True, observed_rows=16),), + workers=(), + to_dict=lambda: {"stage_id": "pretrain", "state": "a" * 64}, + ), + SimpleNamespace( + stage_id="finetune", + state_digest="b" * 64, + roles=(SimpleNamespace(role="train", present=True, observed_rows=12),), + workers=(), + to_dict=lambda: {"stage_id": "finetune", "state": "b" * 64}, + ), + ) + details = _component_state_details(stages) + assert details["component_stage_count"] == 2 + assert details["component_stages"] == "pretrain,finetune" + assert details["anchor_stage"] == "finetune" + assert details["stage.pretrain.rows"] == 16 + assert details["stage.finetune.rows"] == 12 + assert len(details["composition_digest"]) == 64 + + +def test_source_state_details_preserve_adapter_declared_scalars() -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _source_state_details, + ) + + assert _source_state_details( + { + "sampling": "full_neighborhood", + "topology_kind": "relational", + "sparse_routing": "all_to_all_single_owner_mod", + "framework_versions": {"torch": "2.5"}, + } + ) == { + "sampling": "full_neighborhood", + "topology_kind": "relational", + "routing": "all_to_all_single_owner_mod", + "jagged": True, + } + + +def test_replicated_role_evidence_uses_per_rank_rows() -> None: + from tributo.algorithms.api import TorchRoleExecutionEvidence + + evidence = TorchRoleExecutionEvidence( + role="nodes", + mode="replicate", + required=True, + present=True, + empty_rank_policy="reject", + expected_rows=8, + observed_rows=8, + rows_per_rank=(8, 8), + ) + assert evidence.rows_per_rank == (8, 8) + with pytest.raises(AlgorithmConfigurationError): + TorchRoleExecutionEvidence( + role="nodes", + mode="replicate", + required=True, + present=True, + empty_rank_policy="reject", + expected_rows=8, + observed_rows=8, + rows_per_rank=(), + ) + + +def test_adapter_worker_config_cannot_carry_core_paths() -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _validate_adapter_worker_config, + ) + + with pytest.raises(AlgorithmConfigurationError, match="Core-owned path"): + _validate_adapter_worker_config({"ray": {"storage_path": "s3://secret"}}) + + +def test_torch_adapter_context_contains_bindings_but_not_core_control_config() -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _torch_algorithm_context_config, + _torch_input_bindings, + _torch_output_config, + ) + + binding = InputBinding( + name="train", + resolver_id="example.resolver", + reference="memory://train", + feature_names=("feature",), + label_name="label", + ) + plan = SimpleNamespace( + algorithm_config={ + "model": {"width": 4}, + "ray": {"storage_path": "/core/path", "resume": {"uri": "s3://x"}}, + "output": {"bundle_uri": "/core/bundle"}, + }, + input_bindings=SimpleNamespace(bindings=(binding,)), + ) + assert _torch_algorithm_context_config(plan) == {"model": {"width": 4}} + assert _torch_input_bindings(plan)["train"]["feature_names"] == ["feature"] + assert _torch_output_config(plan) == {"bundle_uri": "/core/bundle"} + + +def test_composite_backward_records_reducer_metrics( + monkeypatch: pytest.MonkeyPatch, +) -> None: + torch = pytest.importorskip("torch") + + import tributo.integrations.algorithm_runtimes.ray_train_torch as runtime + + class Reducer: + api_version = 1 + reducer_id = "example.reducer" + component_schema_id = "example.schema" + code_digest = "4" * 64 + + def reduce(self, config, global_state, context): + del config, global_state, context + return TorchGlobalLossReduction( + "accepted", + coefficients={"loss": 1.0}, + metrics={"train_loss": TorchMetricContribution(3.0, 1.0)}, + ) + + reducer = Reducer() + monkeypatch.setattr(runtime, "_load_reference", lambda reference: reducer) + monkeypatch.setattr( + runtime, "_validate_module_digest", lambda reference, digest: None + ) + loss = TorchCompositeLossContribution( + "example.schema", + {"loss": torch.tensor(2.0, requires_grad=True)}, + {"count": 1.0}, + ) + metric_totals: dict[str, list[float]] = {} + value = runtime._composite_backward( + loss, + config={ + "_core_global_loss_reducer_ref": "example.reducer:Reducer", + "_core_composite_loss_schema_id": "example.schema", + "_core_global_loss_reducer_api_version": 1, + "_core_global_loss_reducer_code_digest": reducer.code_digest, + "_core_policy_digest": "1" * 64, + "_core_execution_plan_digest": "2" * 64, + }, + world_size=1, + device=torch.device("cpu"), + dist=torch.distributed, + metric_totals=metric_totals, + ) + assert float(value.detach().item()) == 2.0 + assert metric_totals == {"train_loss": [3.0, 1.0]} + + +def test_component_export_result_contains_composition_digest( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + import tributo.exporting.service as service_module + import tributo.integrations.algorithm_runtimes.ray_train_torch as runtime + import tributo.integrations.sources.ray_torch as source_module + from tributo.exporting.models import ExportSource + + class FakeProvider: + def open_source(self, result, options): + del result, options + + @contextmanager + def opened(): + yield ExportSource( + source_kind="torch_module", + metadata={ + "artifact_plan": { + "targets": ({"name": "model", "format": "onnx"},), + "roles": {"inference": "model"}, + } + }, + ) + + return opened() + + class FakeService: + def export_bundle(self, source, config, *, tributo_version): + del source, config, tributo_version + return SimpleNamespace( + bundle_id="bundle-1", + canonical_uri=str(tmp_path / "bundle"), + execution_id="execution-1", + manifest_sha256="5" * 64, + ) + + monkeypatch.setattr(source_module, "RayTorchSourceProvider", FakeProvider) + monkeypatch.setattr(service_module, "BundleExportService", FakeService) + policy = SimpleNamespace( + loop_owner="core_recipe", + digest="1" * 64, + state_layout="component", + execution_plan=SimpleNamespace( + digest="2" * 64, + stages=(SimpleNamespace(input_roles=("train",)),), + ), + ) + binding = InputBinding( + name="train", + resolver_id="example.resolver", + reference="memory://train", + feature_names=("feature",), + label_name="label", + ) + plan = SimpleNamespace( + plan_id="2" * 64, + distribution_spec=SimpleNamespace( + strategy=DistributionStrategy.RAY_TRAIN_TORCH, + result_policy=ResultPolicy.BUNDLE_REQUIRED, + policy=policy, + ), + algorithm_config={"output": {"bundle_uri": str(tmp_path / "bundle")}}, + implementation=SimpleNamespace( + implementation_ref="example:Implementation", + code_digest="3" * 64, + implementation_id="example.implementation", + ), + input_bindings=SimpleNamespace(bindings=(binding,)), + ) + result = SimpleNamespace( + checkpoint=object(), + metrics={"torch_evidence": {"composition_digest": "a" * 64}}, + core_evidence_attested=True, + ) + execution = runtime.export_ray_train_torch_result( + result=result, + plan=plan, + run_id="aabbccdd", + ) + assert execution.outputs["composition_digest"] == "a" * 64 + + +def test_checkpoint_report_builds_descriptor_from_payload(tmp_path) -> None: + identity = TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "example", + "example.recipe", + "0" * 64, + "1" * 64, + "2" * 64, + ) + runtime = TorchRuntimeContext( + algorithm_config={}, + implementation_id=identity.implementation_id, + world_size=1, + policy_digest=identity.policy_digest, + execution_plan_digest=identity.execution_plan_digest, + run_identity=identity, + input_binding_digest="3" * 64, + ) + stage = TorchStageContext(runtime, "train", 0, True, ("train",)) + payload = tmp_path / "model.pt" + payload.write_bytes(b"model") + captured: dict[str, object] = {} + + class Draft: + checkpoint_dir = tmp_path + + def report(self, *, metrics, stage_context, completed_step) -> None: + captured.update(metrics) + assert stage_context is stage + assert completed_step == 1 + + report_torch_checkpoint( + {"train_loss": 0.5}, + Draft(), + stage, + 1, + ) + descriptor = TorchCheckpointDescriptor.from_dict(captured["checkpoint_descriptor"]) + assert descriptor.identity == identity + assert descriptor.payload_files == { + "model.pt": __import__("hashlib").sha256(b"model").hexdigest() + } + assert (tmp_path / "torch_checkpoint_descriptor.json").is_file() + + +def test_recovery_envelope_roundtrip_and_locator_digest_binding() -> None: + locator = TorchCheckpointLocator("s3://bucket/stage", "4" * 64) + envelope = TorchRecoveryEnvelope( + completed_stage_ids=("pretrain",), + stage_checkpoints={"pretrain": locator}, + active_stage_id="finetune", + active_checkpoint=TorchCheckpointLocator("s3://bucket/active", "5" * 64), + ) + restored = TorchRecoveryEnvelope.from_dict(envelope.to_dict()) + assert restored == envelope + with pytest.raises(AlgorithmConfigurationError): + TorchRecoveryEnvelope( + completed_stage_ids=("pretrain",), + stage_checkpoints={ + "pretrain": TorchCheckpointLocator("s3://bucket/stage", "4" * 64) + }, + active_stage_id="pretrain", + active_checkpoint=TorchCheckpointLocator("s3://bucket/active", "5" * 64), + ) + with pytest.raises(AlgorithmConfigurationError): + TorchRecoveryEnvelope.from_dict( + { + "completed_stage_ids": ["pretrain"], + "stage_checkpoints": {"pretrain": "not-a-locator"}, + } + ) + + +def test_checkpoint_payload_rejects_symlinked_descriptor(tmp_path) -> None: + identity = TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "example", + "example.recipe", + "0" * 64, + "1" * 64, + "2" * 64, + ) + runtime = TorchRuntimeContext( + algorithm_config={}, + implementation_id=identity.implementation_id, + world_size=1, + policy_digest=identity.policy_digest, + execution_plan_digest=identity.execution_plan_digest, + run_identity=identity, + input_binding_digest="3" * 64, + ) + stage = TorchStageContext(runtime, "train", 0, True, ("train",)) + (tmp_path / "model.pt").write_bytes(b"model") + target = tmp_path / "outside.json" + target.write_text("{}", encoding="utf-8") + (tmp_path / "torch_checkpoint_descriptor.json").symlink_to(target) + + class Draft: + checkpoint_dir = tmp_path + + def report(self, *, metrics, stage_context, completed_step) -> None: + del metrics, stage_context, completed_step + + with pytest.raises(AlgorithmExecutionError): + report_torch_checkpoint({}, Draft(), stage, 1) + + +def test_checkpoint_report_rejects_core_metadata_fields(tmp_path) -> None: + identity = TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "example", + "example.recipe", + "0" * 64, + "1" * 64, + "2" * 64, + ) + runtime = TorchRuntimeContext( + algorithm_config={}, + implementation_id=identity.implementation_id, + world_size=1, + policy_digest=identity.policy_digest, + execution_plan_digest=identity.execution_plan_digest, + run_identity=identity, + input_binding_digest="3" * 64, + ) + stage = TorchStageContext(runtime, "train", 0, True, ("train",)) + (tmp_path / "model.pt").write_bytes(b"model") + + class Draft: + checkpoint_dir = tmp_path + + def report(self, *, metrics, stage_context, completed_step) -> None: + del metrics, stage_context, completed_step + + with pytest.raises(AlgorithmConfigurationError): + report_torch_checkpoint( + {"checkpoint_locator": "s3://bucket/private"}, Draft(), stage, 1 + ) + + +def test_checkpoint_progress_roundtrip_and_conditional_resume_serialization() -> None: + progress = TorchCheckpointProgress( + epoch=2, + micro_batch_cursor=3, + optimizer_step=7, + scheduler_step=2, + accumulation_steps=4, + dataset_cursor_by_rank={"0": 3, "1": 3}, + shuffle_seed=44, + ) + assert TorchCheckpointProgress.from_dict(progress.to_dict()) == progress + identity = TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "example", + "example.recipe", + "0" * 64, + "1" * 64, + "2" * 64, + ) + descriptor = TorchCheckpointDescriptor( + schema_version=1, + identity=identity, + run_config_name=identity.run_config_name, + state_layout="component", + world_size=1, + completed_step=7, + policy_digest=identity.policy_digest, + execution_plan_digest=identity.execution_plan_digest, + input_binding_digest="3" * 64, + implementation_code_digest=identity.implementation_code_digest, + payload_files={"model.pt": "4" * 64}, + resume_supported=False, + same_world_size_resume=None, + ) + assert "same_world_size_resume" not in descriptor.to_dict() + assert TorchCheckpointDescriptor.from_dict(descriptor.to_dict()) == descriptor + + +def test_rank_progress_statistics_and_runtime_context_are_typed() -> None: + statistics = TorchRankProgressStatistics( + rows_processed=4, + coverage_totals={"coverage.positive": 2}, + loss_numerator_total=3.0, + loss_normalizer_total=4.0, + metric_totals={"accuracy": (2.0, 4.0)}, + reducer_observation={"branch": "nnpu_normal"}, + ) + assert TorchRankProgressStatistics.from_dict(statistics.to_dict()) == statistics + with pytest.raises(AlgorithmConfigurationError): + TorchRankProgressStatistics.from_dict({"rows_processed": "four"}) + + runtime = TorchRuntimeContext( + algorithm_config={}, + implementation_id="example.adapter", + world_size=1, + policy_digest="1" * 64, + execution_plan_digest="2" * 64, + resume_supported=False, + same_world_size_resume=None, + ) + payload = runtime.to_dict() + assert "same_world_size_resume" not in payload + restored = TorchStageContext.from_dict( + { + "runtime": payload, + "stage_id": "train", + "stage_index": 0, + "is_final": True, + "input_roles": ["train"], + } + ) + assert restored.runtime.same_world_size_resume is None + + +def test_scheduler_boundary_and_recovery_commit_are_fail_closed(tmp_path) -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _require_checkpoint_commit, + _should_apply_epoch_scheduler, + ) + + assert _should_apply_epoch_scheduler( + restore_same_stage=True, + epoch=1, + restored_epoch=1, + restored_epoch_scheduler_applied=False, + ) + assert not _should_apply_epoch_scheduler( + restore_same_stage=True, + epoch=1, + restored_epoch=1, + restored_epoch_scheduler_applied=True, + ) + identity = TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "example", + "example.recipe", + "0" * 64, + "1" * 64, + "2" * 64, + ) + (tmp_path / "model.pt").write_bytes(b"model") + descriptor = TorchCheckpointDescriptor( + schema_version=1, + identity=identity, + run_config_name=identity.run_config_name, + state_layout="replicated", + world_size=1, + completed_step=1, + policy_digest=identity.policy_digest, + execution_plan_digest=identity.execution_plan_digest, + input_binding_digest="3" * 64, + implementation_code_digest=identity.implementation_code_digest, + payload_files={"model.pt": "4" * 64}, + ) + with pytest.raises(AlgorithmExecutionError, match="commit"): + _require_checkpoint_commit(tmp_path, descriptor) + + +def test_local_stage_staging_ignores_prior_partial_attempt(tmp_path) -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _persist_stage_checkpoint, + ) + + identity = TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "example", + "example.recipe", + "0" * 64, + "1" * 64, + "2" * 64, + ) + source = tmp_path / "source" + source.mkdir() + (source / "model.pt").write_bytes(b"model") + run_root = tmp_path / identity.run_config_name + run_root.mkdir() + + class Checkpoint: + @contextmanager + def as_directory(self): + yield source + + # A previous attempt with the same digest must not block a fresh staging + # attempt; only the committed destination is authoritative. + stale_path = run_root / f".stage_checkpoint.staging-{'4' * 64}-old" + stale_path.mkdir() + locator = _persist_stage_checkpoint( + Checkpoint(), + identity=identity, + storage_path=tmp_path, + descriptor_digest="4" * 64, + ) + assert locator == f"ray://{run_root / 'stage_checkpoint'}" + assert (run_root / "stage_checkpoint" / "torch_stage_commit.json").is_file() + + +def test_composite_zero_global_normalizer_fails_before_reducer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import tributo.integrations.algorithm_runtimes.ray_train_torch as runtime + + torch = pytest.importorskip("torch") + dist = pytest.importorskip("torch.distributed") + invoked = False + + class Reducer: + api_version = 1 + reducer_id = "example.reducer" + component_schema_id = "example.components" + code_digest = "0" * 64 + + def reduce(self, config, global_state, context): + nonlocal invoked + del config, global_state, context + invoked = True + return TorchGlobalLossReduction( + "accepted", + coefficients={"loss": 1.0}, + metrics={"train_loss": TorchMetricContribution(1.0, 1.0)}, + ) + + monkeypatch.setattr(runtime, "_validate_module_digest", lambda *args: None) + monkeypatch.setattr(runtime, "_load_reference", lambda reference: Reducer) + loss = TorchCompositeLossContribution( + "example.components", + {"loss": torch.tensor(1.0, requires_grad=True)}, + {"rows": 0.0}, + ) + with pytest.raises(AlgorithmExecutionError, match="normalizer"): + runtime._reduce_composite_loss( + loss, + config={ + "_core_global_loss_reducer_ref": "example:Reducer", + "_core_global_loss_reducer_api_version": 1, + "_core_global_loss_reducer_code_digest": "0" * 64, + "_core_composite_loss_schema_id": "example.components", + "_core_policy_digest": "1" * 64, + "_core_execution_plan_digest": "2" * 64, + }, + world_size=1, + device=torch.device("cpu"), + dist=dist, + ) + assert not invoked + + +def test_helper_signatures_are_public_and_stable() -> None: + from tributo.algorithms.api.torch_runtime import report_torch_checkpoint + + assert list(inspect.signature(report_torch_checkpoint).parameters) == [ + "metrics", + "payload_draft", + "stage_context", + "completed_step", + ] + + +def test_removed_torch_public_surfaces_are_not_exported() -> None: + import tributo.algorithms as algorithms + + assert not hasattr(algorithms, "TorchTrainingRecipe") + assert not hasattr(algorithms, "TrainingRecipeV2") + assert not hasattr(algorithms.AlgorithmBuilder, "from_torch_recipe") + assert not hasattr(algorithms.AlgorithmBuilder, "from_training_recipe_v2") diff --git a/tests/algorithms/test_training_recipe_v2_contract.py b/tests/algorithms/test_training_recipe_v2_contract.py index 6bd7621..0b23602 100644 --- a/tests/algorithms/test_training_recipe_v2_contract.py +++ b/tests/algorithms/test_training_recipe_v2_contract.py @@ -1,97 +1,19 @@ -"""Contract tests for low-code TrainingRecipeV2.""" +"""Regression tests for removal of the legacy Torch RecipeV2 contract.""" from __future__ import annotations -from tributo.algorithms import AlgorithmBuilder, TrainingRecipeV2 -from tributo.algorithms.api import ( - ContractBinding, - ContractBindingSet, - DistributionStrategy, - EnvironmentSpec, - ExecutionMode, - ExecutionProfile, - MetricReduction, - QualifiedReference, - RuntimeTopology, - WorkerRange, - WorkerResources, -) -from tributo.training.algorithm_spec import AlgorithmSpec +import pytest -from .conftest import make_spec +from tributo.algorithms import AlgorithmBuilder -def _binding(contract_id: str, digest: str, validator: str) -> ContractBinding: - return ContractBinding( - contract_id=contract_id, - schema_version=1, - schema_digest=digest * 64, - validator_ref=QualifiedReference.parse( - f"tests.support.portable_contracts:{validator}" - ), - ) +def test_legacy_recipe_v2_is_not_exported() -> None: + import tributo.algorithms as algorithms + assert not hasattr(algorithms, "TrainingRecipeV2") + assert not hasattr(AlgorithmBuilder, "from_training_recipe_v2") -def _contracts(spec: AlgorithmSpec) -> ContractBindingSet: - assert spec.config_contract_ref is not None - assert spec.input_contract_ref is not None - assert spec.output_contract_ref is not None - return ContractBindingSet( - config=_binding(spec.config_contract_ref, "a", "ConfigValidator"), - input=_binding(spec.input_contract_ref, "b", "InputValidator"), - output=_binding(spec.output_contract_ref, "c", "OutputValidator"), - coverage=_binding("test.recipe.coverage.v1", "d", "CoverageValidator"), - ) - -def test_recipe_v2_requires_only_math_and_codec_hooks() -> None: - assert TrainingRecipeV2.__abstractmethods__ == { - "batch_adapter", - "build_modules", - "checkpoint_codec", - "metric_plan", - "optimization_plan", - "training_step", - "validation_step", - } - - -def test_builder_selects_dedicated_recipe_v2_runtime() -> None: - spec = make_spec( - "external_recipe_v2", - operations=("fit",), - mode=ExecutionMode.TRAINING_RECIPE_V2, - ) - descriptor = AlgorithmBuilder.from_training_recipe_v2( - spec=spec, - implementation_id="example.recipe_v2.binary_linear", - implementation_version="1.0.0", - recipe="tests.support.training_recipe_v2:BinaryLinearRecipeV2", - environment=EnvironmentSpec( - environment_id="example.recipe_v2.v1", - dependencies=("example-recipe-v2==1.0.0", "torch>=2.5"), - ), - metric_reducers={"accuracy": MetricReduction.SUM_COUNT}, - supported_worker_range=WorkerRange(1, 8), - supported_execution_profiles=( - ExecutionProfile.LOCAL, - ExecutionProfile.CLUSTER, - ), - resources_per_worker=WorkerResources(num_cpus=1), - package_name="example-recipe-v2", - package_version="1.0.0", - tributo_version_spec=">=1,<2", - contract_bindings=_contracts(spec), - ) - - registration = descriptor.registration - assert registration.distribution_spec is not None - assert ( - registration.distribution_spec.strategy - is DistributionStrategy.RAY_TRAIN_RECIPE_V2 - ) - assert registration.implementation.runtime_id == "tributo.ray_train_recipe_v2" - assert registration.implementation.input_compatibility.distribution_policy == ( - RuntimeTopology.RAY_TRAIN_RECIPE_V2, - ) - assert descriptor.api_version == 2 +def test_legacy_recipe_v2_module_is_not_importable() -> None: + with pytest.raises(ImportError): + __import__("tributo.integrations.algorithm_runtimes.torch_recipe") diff --git a/tests/fixtures/torch_recipe_algorithm_plugin/src/tributo_test_torch_recipe_algorithm/__init__.py b/tests/fixtures/torch_recipe_algorithm_plugin/src/tributo_test_torch_recipe_algorithm/__init__.py index 4de9e4b..13b7896 100644 --- a/tests/fixtures/torch_recipe_algorithm_plugin/src/tributo_test_torch_recipe_algorithm/__init__.py +++ b/tests/fixtures/torch_recipe_algorithm_plugin/src/tributo_test_torch_recipe_algorithm/__init__.py @@ -2,10 +2,22 @@ from __future__ import annotations +import hashlib from collections.abc import Mapping -from typing import Any - -from tributo.algorithms import AlgorithmBuilder, TorchTrainingRecipe +from pathlib import Path + +from tributo.algorithms import ( + AlgorithmBuilder, + TorchArtifactPlan, + TorchBatch, + TorchBatchContext, + TorchLossContribution, + TorchMetricPlan, + TorchModuleSet, + TorchOptimizationPlan, + TorchRecipe, + TorchStepResult, +) from tributo.algorithms.api import ( EnvironmentSpec, ExecutionMode, @@ -20,51 +32,96 @@ ProblemType, ) +CODE_DIGEST = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + -class ThirdPartyBinaryLinearRecipe(TorchTrainingRecipe): - """Define only the four model-level factories required by the recipe SPI.""" +class ThirdPartyBinaryLinearRecipe(TorchRecipe): + """A low-dependency out-of-tree implementation of the new Recipe SPI.""" - def model_factory(self, config: Mapping[str, Any]) -> object: + def build_modules(self, context: object) -> TorchModuleSet: import torch - return torch.nn.Linear(int(config.get("input_features", 2)), 1) + return TorchModuleSet( + {"model": torch.nn.Linear(2, 1), "loss": torch.nn.BCEWithLogitsLoss()} + ) - def loss_factory(self, config: Mapping[str, Any]) -> object: + def adapt_batch(self, batch: object, context: object) -> TorchBatch: import torch - del config - return torch.nn.BCEWithLogitsLoss() + if not isinstance(batch, Mapping): + raise TypeError("batch must be a mapping") + if not isinstance(context, TorchBatchContext): + raise TypeError("Torch batch context is invalid") + features = torch.stack( + [ + torch.as_tensor(batch[name], dtype=torch.float32) + for name in context.feature_names + ], + dim=1, + ) + if context.label_name is None: + raise ValueError("Torch fixture requires a label") + targets = torch.as_tensor( + batch[context.label_name], dtype=torch.float32 + ).reshape(-1, 1) + return TorchBatch( + positional=(features,), targets=targets, local_rows=len(targets) + ) - def optimizer_factory( - self, - model: object, - config: Mapping[str, Any], - ) -> object: + def training_step( + self, modules: TorchModuleSet, batch: TorchBatch, context: object + ) -> TorchStepResult: import torch - if not isinstance(model, torch.nn.Module): - raise TypeError("model must be torch.nn.Module") - return torch.optim.SGD( - model.parameters(), - lr=float(config.get("learning_rate", 0.1)), + del context + predictions = modules["model"](batch.positional[0]) + numerator = torch.nn.functional.binary_cross_entropy_with_logits( + predictions, batch.targets.float(), reduction="sum" + ) + return TorchStepResult( + outputs={"prediction": predictions}, + loss=TorchLossContribution(numerator, batch.local_rows), ) - def metric_factories(self, config: Mapping[str, Any]) -> Mapping[str, Any]: - del config + def validation_step( + self, modules: TorchModuleSet, batch: TorchBatch, context: object + ) -> TorchStepResult: + return self.training_step(modules, batch, context) - def accuracy(predictions: object, targets: object) -> object: - import torch + def configure_optimizers( + self, modules: TorchModuleSet, context: object + ) -> TorchOptimizationPlan: + import torch - if not isinstance(predictions, torch.Tensor) or not isinstance( - targets, torch.Tensor - ): - raise TypeError("accuracy requires Tensor values") - return (torch.sigmoid(predictions) >= 0.5) == targets.bool() + del context + return TorchOptimizationPlan( + torch.optim.SGD(modules["model"].parameters(), lr=0.1) + ) - return {"accuracy": accuracy} + def metric_plan(self, context: object) -> TorchMetricPlan: + del context + return TorchMetricPlan({"accuracy": "sum_count", "train_loss": "sum_count"}) + + def artifact_plan(self, context: object) -> TorchArtifactPlan: + del context + return TorchArtifactPlan( + source_kind="torch_module", + input_signature=( + { + "name": "features", + "dtype": "float32", + "shape": ("batch", 2), + }, + ), + output_signature=( + {"name": "prediction", "dtype": "float32", "shape": ("batch", 1)}, + ), + targets=({"name": "onnx-model", "format": "onnx"},), + roles={"inference": "onnx-model"}, + ) -DESCRIPTOR = AlgorithmBuilder.from_torch_recipe( +DESCRIPTOR = AlgorithmBuilder.from_torch( spec=AlgorithmSpec( name="third_party_binary_linear", trainer_cls=None, @@ -79,7 +136,7 @@ def accuracy(predictions: object, targets: object) -> object: model_family="linear_model", data_modalities=("tabular",), lifecycle_kind="batch_fit", - allowed_execution_modes=(ExecutionMode.COLLECTIVE.value,), + allowed_execution_modes=(ExecutionMode.RAY_TRAIN_TORCH.value,), config_contract_ref="example.torch_recipe.config.v1", input_contract_ref="tributo.tabular.dense.v1", output_contract_ref="tributo.classification.onnx.v1", @@ -105,6 +162,8 @@ def accuracy(predictions: object, targets: object) -> object: package_name="tributo-test-torch-recipe-algorithm", package_version="0.1.0", tributo_version_spec=">=1,<2", + code_digest=CODE_DIGEST, + descriptor_api_version=1, stability="alpha", tested=True, supported=True, diff --git a/tests/inference/test_contracts.py b/tests/inference/test_contracts.py index b78d871..36edda4 100644 --- a/tests/inference/test_contracts.py +++ b/tests/inference/test_contracts.py @@ -2,6 +2,7 @@ from __future__ import annotations +import numpy as np import pytest from pydantic import ValidationError @@ -27,6 +28,7 @@ TensorInputBinding, TensorOutputBinding, ) +from tributo.inference.kernel import _build_input_tensor def _request(**updates) -> InferenceRequest: @@ -252,6 +254,54 @@ def test_single_column_mode_is_explicit_and_json_stable(self) -> None: TensorInputBinding.model_validate_json(scalar.model_dump_json()) == scalar ) + def test_single_vector_column_preserves_nested_tensor_rank(self) -> None: + column = np.empty(2, dtype=object) + column[0] = [[1.0], [2.0]] + column[1] = [[3.0], [4.0]] + tensor = _build_input_tensor( + {"window": column}, + columns=("window",), + dtype="float32", + single_column_mode="vector", + null_policy="error", + nan_policy="error", + ) + assert tensor.shape == (2, 2, 1) + + def test_multi_dimensional_object_column_preserves_nested_tensor_rank( + self, + ) -> None: + column = np.empty((2, 2), dtype=object) + column[0] = [np.asarray([1.0, 2.0]), np.asarray([3.0, 4.0])] + column[1] = [np.asarray([5.0, 6.0]), np.asarray([7.0, 8.0])] + tensor = _build_input_tensor( + {"window": column}, + columns=("window",), + dtype="float32", + single_column_mode="vector", + null_policy="error", + nan_policy="error", + ) + assert tensor.shape == (2, 2, 2) + + def test_arrow_nested_object_rows_preserve_tensor_rank(self) -> None: + column = np.empty(2, dtype=object) + column[0] = np.asarray( + [np.asarray([1.0, 2.0]), np.asarray([3.0, 4.0])], dtype=object + ) + column[1] = np.asarray( + [np.asarray([5.0, 6.0]), np.asarray([7.0, 8.0])], dtype=object + ) + tensor = _build_input_tensor( + {"window": column}, + columns=("window",), + dtype="float32", + single_column_mode="vector", + null_policy="error", + nan_policy="error", + ) + assert tensor.shape == (2, 2, 2) + def test_scalar_single_column_mode_rejects_invalid_contracts(self) -> None: with pytest.raises(ValidationError, match="requires exactly one column"): TensorInputBinding( diff --git a/tests/integration/test_distributed_algorithm_it_contract.py b/tests/integration/test_distributed_algorithm_it_contract.py index 9cc059b..fa8323e 100644 --- a/tests/integration/test_distributed_algorithm_it_contract.py +++ b/tests/integration/test_distributed_algorithm_it_contract.py @@ -561,8 +561,8 @@ def test_torch_recipe_fixture_is_a_code_only_low_code_wheel() -> None: assert set(pyproject["project"]["entry-points"]["tributo.algorithms"]) == { "third_party_binary_linear" } - assert "AlgorithmBuilder.from_torch_recipe" in fixture_source - assert "TorchTrainingRecipe" in fixture_source + assert "AlgorithmBuilder.from_torch" in fixture_source + assert "TorchRecipe" in fixture_source assert "train_loop_per_worker" not in fixture_source assert "tributo.algorithms.builtin" not in fixture_source diff --git a/tests/support/torch_recipe.py b/tests/support/torch_recipe.py index 77c9b45..661b718 100644 --- a/tests/support/torch_recipe.py +++ b/tests/support/torch_recipe.py @@ -1,54 +1,107 @@ -"""Trusted test recipe that imports PyTorch only when its factories run.""" +"""Trusted test TorchRecipe that imports PyTorch only inside hooks.""" from __future__ import annotations from collections.abc import Mapping -from typing import Any -from tributo.algorithms import TorchTrainingRecipe +from tributo.algorithms import ( + TorchArtifactPlan, + TorchBatch, + TorchBatchContext, + TorchLossContribution, + TorchMetricPlan, + TorchModuleSet, + TorchOptimizationPlan, + TorchRecipe, + TorchStepResult, +) -class BinaryLinearRecipe(TorchTrainingRecipe): - """Minimal dense binary classifier used by contract and integration tests.""" +class BinaryLinearRecipe(TorchRecipe): + """Minimal dense binary classifier used by Core contract tests.""" - def model_factory(self, config: Mapping[str, Any]) -> object: + def build_modules(self, context: object) -> TorchModuleSet: import torch - return torch.nn.Linear(int(config.get("input_features", 2)), 1) + return TorchModuleSet( + {"model": torch.nn.Linear(2, 1), "loss": torch.nn.BCEWithLogitsLoss()} + ) - def loss_factory(self, config: Mapping[str, Any]) -> object: + def adapt_batch(self, batch: object, context: object) -> TorchBatch: import torch - del config - return torch.nn.BCEWithLogitsLoss() + if not isinstance(batch, Mapping): + raise TypeError("Torch test batch must be a mapping") + if not isinstance(context, TorchBatchContext): + raise TypeError("Torch test batch context is invalid") + features = torch.stack( + [batch[name] for name in context.feature_names], + dim=1, + ) + if context.label_name is None: + raise ValueError("Torch test recipe requires a label") + targets = batch[context.label_name].reshape(-1, 1) + return TorchBatch( + positional=(features,), targets=targets, local_rows=len(targets) + ) - def optimizer_factory( - self, - model: object, - config: Mapping[str, Any], - ) -> object: + def training_step( + self, modules: TorchModuleSet, batch: TorchBatch, context: object + ) -> TorchStepResult: import torch - if not isinstance(model, torch.nn.Module): - raise TypeError("model must be torch.nn.Module") - return torch.optim.SGD( - model.parameters(), - lr=float(config.get("learning_rate", 0.1)), + del context + predictions = modules["model"](batch.positional[0]) + numerator = torch.nn.functional.binary_cross_entropy_with_logits( + predictions, batch.targets.float(), reduction="sum" + ) + return TorchStepResult( + outputs={"prediction": predictions}, + loss=TorchLossContribution(numerator, batch.local_rows), ) - def metric_factories(self, config: Mapping[str, Any]) -> Mapping[str, Any]: - del config + def validation_step( + self, modules: TorchModuleSet, batch: TorchBatch, context: object + ) -> TorchStepResult: + return self.training_step(modules, batch, context) - def accuracy(predictions: object, targets: object) -> object: - import torch + def configure_optimizers( + self, modules: TorchModuleSet, context: object + ) -> TorchOptimizationPlan: + import torch - if not isinstance(predictions, torch.Tensor) or not isinstance( - targets, torch.Tensor - ): - raise TypeError("accuracy requires Tensor values") - return (torch.sigmoid(predictions) >= 0.5) == targets.bool() + del context + return TorchOptimizationPlan( + torch.optim.SGD(modules["model"].parameters(), lr=0.1) + ) - return {"accuracy": accuracy} + def metric_plan(self, context: object) -> TorchMetricPlan: + del context + return TorchMetricPlan({"accuracy": "sum_count", "train_loss": "sum_count"}) + + def artifact_plan(self, context: object) -> TorchArtifactPlan: + del context + return TorchArtifactPlan( + source_kind="torch_module", + input_signature=( + { + "name": "features", + "dtype": "float32", + "shape": ("batch", 2), + }, + ), + output_signature=( + {"name": "prediction", "dtype": "float32", "shape": ("batch", 1)}, + ), + targets=( + { + "name": "onnx-model", + "format": "onnx", + "exporter_id": "torch-onnx-v1", + }, + ), + roles={"inference": "onnx-model"}, + ) __all__ = ["BinaryLinearRecipe"] diff --git a/tests/support/training_recipe_v2.py b/tests/support/training_recipe_v2.py index b7f05c3..8e589e0 100644 --- a/tests/support/training_recipe_v2.py +++ b/tests/support/training_recipe_v2.py @@ -1,4 +1,4 @@ -"""Independent TrainingRecipeV2 fixture without Ray orchestration code.""" +"""Migrated TorchRecipe fixture kept for compatibility-focused test imports.""" from __future__ import annotations @@ -7,10 +7,19 @@ from typing import Any, cast from tributo.algorithms import ( - MetricPlan, - OptimizationPlan, - TrainingRecipeV2, - TrainingStepResult, + TorchArtifactContext, + TorchArtifactPlan, + TorchBatch, + TorchBatchContext, + TorchBuildContext, + TorchLossContribution, + TorchMetricContribution, + TorchMetricPlan, + TorchModuleSet, + TorchOptimizationPlan, + TorchRecipe, + TorchStepContext, + TorchStepResult, ) @@ -23,7 +32,9 @@ def _accuracy(predictions: object, targets: object) -> object: return (predicted == target_tensor).to(dtype=torch.float32).mean() -class PickleCheckpointCodec: +class CheckpointCodec: + """Legacy fixture codec retained only for migration tests.""" + def dumps(self, value: object) -> bytes: return pickle.dumps(value, protocol=5) @@ -31,93 +42,115 @@ def loads(self, payload: bytes) -> object: return pickle.loads(payload) -class BinaryLinearRecipeV2(TrainingRecipeV2): - """Define only modules, batch conversion, Step, Plan, and Codec.""" +class BinaryLinearRecipeV2(TorchRecipe): + """The old fixture name using the new typed TorchRecipe contract.""" - def build_modules(self, config: Mapping[str, Any]) -> Mapping[str, object]: + def build_modules(self, context: TorchBuildContext) -> TorchModuleSet: import torch - model_config = config.get("model", {}) + model_config = context.runtime.algorithm_config.get("model", {}) if not isinstance(model_config, Mapping): raise ValueError("model config must be a mapping") feature_count = int(model_config.get("input_features", 2)) - return { - "model": torch.nn.Linear(feature_count, 1), - "loss": torch.nn.BCEWithLogitsLoss(), - } + return TorchModuleSet( + { + "model": torch.nn.Linear(feature_count, 1), + "loss": torch.nn.BCEWithLogitsLoss(), + } + ) - def batch_adapter( + def adapt_batch( self, batch: object, - *, - feature_names: tuple[str, ...], - label_name: str | None, - weight_name: str | None, - config: Mapping[str, Any], - ) -> tuple[object, object, object | None, int]: + context: TorchBatchContext, + ) -> TorchBatch: import torch - del config if not isinstance(batch, Mapping): raise ValueError("batch must be columnar") features = torch.stack( - [batch[name].to(dtype=torch.float32) for name in feature_names], + [batch[name].to(dtype=torch.float32) for name in context.feature_names], dim=1, ) - if label_name is None: + if context.label_name is None: raise ValueError("binary fixture requires a label") - targets = batch[label_name].to(dtype=torch.float32).reshape(-1, 1) - weights = batch.get(weight_name) if weight_name is not None else None - return features, targets, weights, int(features.shape[0]) + targets = batch[context.label_name].to(dtype=torch.float32).reshape(-1, 1) + weights = ( + batch.get(context.weight_name) if context.weight_name is not None else None + ) + return TorchBatch( + positional=(features,), + targets=targets, + weights=weights, + local_rows=int(features.shape[0]), + ) def training_step( self, - modules: Mapping[str, object], - features: object, - targets: object, - weights: object | None, - config: Mapping[str, Any], - ) -> TrainingStepResult: - del weights, config + modules: TorchModuleSet, + batch: TorchBatch, + context: TorchStepContext, + ) -> TorchStepResult: + del context model = cast(Any, modules["model"]) loss = cast(Any, modules["loss"]) - predictions = model(features) - return TrainingStepResult( - predictions=predictions, loss=loss(predictions, targets) + predictions = model(batch.positional[0]) + loss_numerator = loss(predictions, batch.targets) * batch.local_rows + accuracy = _accuracy(predictions, batch.targets) + return TorchStepResult( + outputs={"prediction": predictions}, + loss=TorchLossContribution(loss_numerator, batch.local_rows), + metrics={ + "accuracy": TorchMetricContribution( + float(accuracy.detach().item()) * batch.local_rows, + batch.local_rows, + ) + }, ) def validation_step( self, - modules: Mapping[str, object], - features: object, - targets: object, - weights: object | None, - config: Mapping[str, Any], - ) -> TrainingStepResult: - return self.training_step(modules, features, targets, weights, config) - - def optimization_plan( + modules: TorchModuleSet, + batch: TorchBatch, + context: TorchStepContext, + ) -> TorchStepResult: + return self.training_step(modules, batch, context) + + def configure_optimizers( self, - model: object, - config: Mapping[str, Any], - ) -> OptimizationPlan: + modules: TorchModuleSet, + context: TorchBuildContext, + ) -> TorchOptimizationPlan: import torch - return OptimizationPlan( + return TorchOptimizationPlan( optimizer=torch.optim.SGD( - cast(Any, model).parameters(), - lr=float(config.get("learning_rate", 0.1)), + cast(Any, modules["model"]).parameters(), + lr=float(context.runtime.algorithm_config.get("learning_rate", 0.1)), + ), + gradient_accumulation_steps=int( + context.runtime.algorithm_config.get("accumulation_steps", 1) + ), + max_gradient_norm=float( + context.runtime.algorithm_config.get("max_gradient_norm", 1.0) ), - gradient_accumulation_steps=int(config.get("accumulation_steps", 1)), - max_gradient_norm=float(config.get("max_gradient_norm", 1.0)), ) - def metric_plan(self, config: Mapping[str, Any]) -> MetricPlan: - del config - return MetricPlan(factories={"accuracy": _accuracy}) - - def checkpoint_codec(self) -> object: - return PickleCheckpointCodec() + def metric_plan(self, context: object) -> TorchMetricPlan: + del context + return TorchMetricPlan({"accuracy": "sum_count", "train_loss": "sum_count"}) + + def artifact_plan(self, context: TorchArtifactContext) -> TorchArtifactPlan: + del context + return TorchArtifactPlan( + source_kind="torch_module", + input_signature=(), + output_signature=( + {"name": "prediction", "dtype": "float32", "shape": ("batch", 1)}, + ), + targets=(), + roles={}, + ) -__all__ = ["BinaryLinearRecipeV2", "PickleCheckpointCodec"] +__all__ = ["BinaryLinearRecipeV2", "CheckpointCodec"] diff --git a/tests/test_plugin.py b/tests/test_plugin.py index 24d3999..1e2ec00 100644 --- a/tests/test_plugin.py +++ b/tests/test_plugin.py @@ -419,8 +419,8 @@ def package_version(name: str) -> str: assert diagnostics == [] assert ( str(fixture.DESCRIPTOR.registration.implementation.executable_factory_ref) - == "tributo.integrations.algorithm_runtimes.torch_recipe:" - "create_torch_recipe_algorithm" + == "tributo.integrations.algorithm_runtimes.ray_train_torch:" + "create_torch_algorithm" ) @@ -535,7 +535,7 @@ def package_version(name: str) -> str: monkeypatch.setattr(importlib.metadata, "version", package_version) diagnostics = [] - with pytest.raises(TypeError, match="TorchTrainingRecipe"): + with pytest.raises(TypeError, match="TorchRecipe"): plugin.validate_distributed_algorithm_descriptor( invalid_descriptor, entry_point_name=descriptor.name, diff --git a/tests/test_stability_inventory.py b/tests/test_stability_inventory.py index a229a76..68754e8 100644 --- a/tests/test_stability_inventory.py +++ b/tests/test_stability_inventory.py @@ -44,6 +44,7 @@ "tributo.algorithms.api.execution": "alpha", "tributo.algorithms.api.models": "alpha", "tributo.algorithms.api.support": "alpha", + "tributo.algorithms.api.torch_runtime": "alpha", "tributo.algorithms.conformance": "alpha", "tributo.algorithms.composition": "alpha", "tributo.algorithms.core.builder": "alpha", @@ -153,7 +154,7 @@ "tributo.integrations.model_runtimes.resolver": "alpha", "tributo.integrations.validators": "beta", "tributo.integrations.sources": "beta", - "tributo.integrations.sources.ray_torch_recipe": "alpha", + "tributo.integrations.sources.ray_torch": "alpha", "tributo.integrations.storage": "beta", "tributo.integrations.hooks": "beta", "tributo.integrations.sinks.parquet": "alpha", diff --git a/tests/training/exporters/test_architecture_contracts.py b/tests/training/exporters/test_architecture_contracts.py index d58585e..a61509a 100644 --- a/tests/training/exporters/test_architecture_contracts.py +++ b/tests/training/exporters/test_architecture_contracts.py @@ -209,7 +209,7 @@ def test_first_party_source_providers_have_stable_unique_ids() -> None: providers = first_party_source_providers() assert {provider.provider_id for provider in providers} == { - "ray-torch-recipe-v1", + "ray-torch-v1", } assert len({provider.trainer_type for provider in providers}) == len(providers) diff --git a/tests/training/exporters/test_torch_recipe_source.py b/tests/training/exporters/test_torch_recipe_source.py index 8c79b13..8b4076c 100644 --- a/tests/training/exporters/test_torch_recipe_source.py +++ b/tests/training/exporters/test_torch_recipe_source.py @@ -2,78 +2,275 @@ from __future__ import annotations +import hashlib import json +from contextlib import contextmanager from pathlib import Path import pytest -from tests.support.torch_recipe import BinaryLinearRecipe -from tributo.exporting.models import BundleOutputConfig, ExportTarget +from tributo.algorithms import ( + TorchArtifactPlan, + TorchCheckpointDescriptor, + TorchCheckpointRef, + TorchLossContribution, + TorchMetricPlan, + TorchModuleSet, + TorchOptimizationPlan, + TorchRecipe, + TorchRuntimeContext, + TorchStageContext, + TorchStageRunIdentity, + TorchStepResult, +) +from tributo.algorithms.spi import ( + RayTorchAdapter, + TorchArtifactContext, + TorchCheckpointContext, + TorchWorkerCheckpointContext, +) +from tributo.exporting.models import BundleOutputConfig, ExportSource, ExportTarget from tributo.exporting.service import BundleExportService -from tributo.integrations.algorithm_runtimes.torch_recipe import _checkpoint_contract -from tributo.integrations.sources.ray_torch_recipe import ( - RayTorchRecipeSourceProvider, - TorchRecipeSourceOptions, +from tributo.integrations.sources.ray_torch import ( + RayTorchSourceProvider, + TorchSourceOptions, ) torch = pytest.importorskip("torch") +CODE_DIGEST = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + + +class BinaryLinearRecipe(TorchRecipe): + def build_modules(self, context): + return TorchModuleSet( + {"model": torch.nn.Linear(2, 1), "loss": torch.nn.MSELoss()} + ) + + def adapt_batch(self, batch, context): + features = torch.as_tensor(batch["features"]) + targets = torch.as_tensor(batch["label"]) + return __import__("tributo.algorithms", fromlist=["TorchBatch"]).TorchBatch( + positional=(features,), targets=targets, local_rows=len(targets) + ) + + def training_step(self, modules, batch, context): + predictions = modules["model"](batch.positional[0]) + loss = torch.nn.functional.mse_loss(predictions, batch.targets, reduction="sum") + return TorchStepResult( + outputs={"prediction": predictions}, + loss=TorchLossContribution(loss, batch.local_rows), + ) + + def validation_step(self, modules, batch, context): + return self.training_step(modules, batch, context) + + def configure_optimizers(self, modules, context): + return TorchOptimizationPlan( + torch.optim.SGD(modules["model"].parameters(), lr=0.1) + ) + + def metric_plan(self, context): + return __import__( + "tributo.algorithms", fromlist=["TorchMetricPlan"] + ).TorchMetricPlan({"train_loss": "sum_count"}) + + def artifact_plan(self, context): + return TorchArtifactPlan( + source_kind="torch_module", + input_signature=( + { + "name": "features", + "dtype": "float32", + "shape": ("batch", 2), + }, + ), + output_signature=( + {"name": "prediction", "dtype": "float32", "shape": ("batch", 1)}, + ), + targets=( + { + "name": "onnx-model", + "format": "onnx", + "exporter_id": "torch-onnx-v1", + }, + ), + roles={"inference": "onnx-model"}, + ) + + +class AdapterExportFixture(RayTorchAdapter): + """Minimal Adapter proving Core attaches its typed Artifact Plan.""" + + def validate_environment(self, context: TorchRuntimeContext) -> None: + del context + + def bind_datasets(self, datasets, context: TorchStageContext): + del context + return datasets + + def worker_config(self, context: TorchStageContext): + del context + return {} + + def train_loop_per_worker( + self, + worker_config, + checkpoint_context: TorchWorkerCheckpointContext, + ) -> None: + del worker_config, checkpoint_context + + def checkpoint_source( + self, result: object, context: TorchCheckpointContext + ) -> object: + del context + return result.checkpoint + + def metric_plan(self, context: TorchRuntimeContext) -> TorchMetricPlan: + del context + return TorchMetricPlan({"train_loss": "sum_count"}) + + def artifact_plan(self, context: TorchArtifactContext) -> TorchArtifactPlan: + if set(context.stage.runtime.algorithm_config) & {"ray", "output"}: + raise AssertionError("Core control config leaked into Adapter context") + if "train" not in context.stage.runtime.input_bindings: + raise AssertionError( + "InputBinding metadata is missing from Adapter context" + ) + return TorchArtifactPlan( + source_kind="torch_module", + input_signature=( + {"name": "features", "dtype": "float32", "shape": ("batch", 2)}, + ), + output_signature=( + {"name": "prediction", "dtype": "float32", "shape": ("batch", 1)}, + ), + targets=({"name": "model", "format": "onnx"},), + roles={"inference": "model"}, + ) + + @contextmanager + def open_export_source( + self, + checkpoint_ref: TorchCheckpointRef, + artifact_context: TorchArtifactContext, + ): + del checkpoint_ref, artifact_context + yield ExportSource( + source_kind="torch_module", + model_object=torch.nn.Linear(2, 1), + sample_inputs={"features": torch.zeros((1, 2))}, + ) + def _checkpoint(path: Path) -> Path: torch.manual_seed(7) - model = BinaryLinearRecipe().model_factory({"input_features": 2}) + model = BinaryLinearRecipe().build_modules(None)["model"] assert isinstance(model, torch.nn.Module) torch.save(model.state_dict(), path / "model.pt") - payload = _checkpoint_contract( - config={ - "model": {"input_features": 2}, - "_tributo_implementation_id": "example.binary_linear", - "_tributo_algorithm": "binary_linear", - "_tributo_feature_names": ["x1", "x2"], - "_tributo_recipe_ref": ("tests.support.torch_recipe:BinaryLinearRecipe"), - "_tributo_recipe_code_digest": None, + identity = TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "binary", + "example.binary", + CODE_DIGEST, + "1" * 64, + "2" * 64, + plan_digest="2" * 64, + ) + descriptor = TorchCheckpointDescriptor( + schema_version=1, + identity=identity, + run_config_name=identity.run_config_name, + state_layout="replicated", + world_size=1, + completed_step=1, + policy_digest=identity.policy_digest, + execution_plan_digest=identity.execution_plan_digest, + input_binding_digest="3" * 64, + implementation_code_digest=CODE_DIGEST, + payload_files={ + "model.pt": hashlib.sha256((path / "model.pt").read_bytes()).hexdigest() }, - feature_count=2, - output_shape=(1,), - framework_version=torch.__version__, ) - (path / "model_config.json").write_text(json.dumps(payload), encoding="utf-8") + (path / "torch_checkpoint_descriptor.json").write_text( + json.dumps(descriptor.to_dict()), encoding="utf-8" + ) (path / "metrics.json").write_text( json.dumps({"train_loss": 0.5}), encoding="utf-8" ) return path -def _options() -> TorchRecipeSourceOptions: - return TorchRecipeSourceOptions( - recipe_ref="tests.support.torch_recipe:BinaryLinearRecipe", - recipe_code_digest=None, - implementation_id="example.binary_linear", +def _adapter_checkpoint(path: Path) -> Path: + (path / "model.pt").write_bytes(b"adapter-model") + identity = TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "adapter", + "example.adapter", + CODE_DIGEST, + "1" * 64, + "2" * 64, + plan_digest="2" * 64, + ) + descriptor = TorchCheckpointDescriptor( + schema_version=1, + identity=identity, + run_config_name=identity.run_config_name, + state_layout="component", + world_size=1, + completed_step=1, + policy_digest=identity.policy_digest, + execution_plan_digest=identity.execution_plan_digest, + input_binding_digest="3" * 64, + implementation_code_digest=CODE_DIGEST, + payload_files={ + "model.pt": hashlib.sha256((path / "model.pt").read_bytes()).hexdigest() + }, + adapter_identity=identity.implementation_id, + resume_supported=False, + same_world_size_resume=None, + ) + (path / "torch_checkpoint_descriptor.json").write_text( + json.dumps(descriptor.to_dict()), encoding="utf-8" + ) + return path + + +def _options() -> TorchSourceOptions: + return TorchSourceOptions( + implementation_ref="tests.training.exporters.test_torch_recipe_source:BinaryLinearRecipe", + implementation_code_digest=CODE_DIGEST, + implementation_id="example.binary", + policy_digest="1" * 64, + plan_digest="2" * 64, + input_binding_digest="3" * 64, + algorithm_config={"model": {"input_features": 2}}, ) def test_recipe_source_reconstructs_model_and_checkpoint_contract( tmp_path: Path, ) -> None: - provider = RayTorchRecipeSourceProvider() + provider = RayTorchSourceProvider() with provider.open_source(_checkpoint(tmp_path), _options()) as source: assert isinstance(source.model_object, torch.nn.Module) assert source.source_kind == "torch_module" - assert source.architecture_id == "example.binary_linear" - assert source.model_config_data == {"input_features": 2} - assert set(source.sample_inputs) == {"x1", "x2"} - assert source.sample_inputs["x1"].shape == (2,) - assert source.feature_schema["feature_names"] == ["x1", "x2"] - assert source.checkpoint_contract is not None - assert source.checkpoint_contract.trainer_type == "torch_recipe" + assert source.architecture_id == "example.binary" + assert source.metadata["artifact_plan"]["roles"] == {"inference": "onnx-model"} def test_recipe_source_rejects_plan_identity_drift(tmp_path: Path) -> None: - provider = RayTorchRecipeSourceProvider() + provider = RayTorchSourceProvider() options = _options().model_copy(update={"implementation_id": "other.model"}) - with pytest.raises(ValueError, match="implementation identity"): + with pytest.raises(Exception, match="identity"): with provider.open_source(_checkpoint(tmp_path), options): pass @@ -82,7 +279,7 @@ def test_recipe_source_uses_existing_onnx_bundle_pipeline(tmp_path: Path) -> Non checkpoint_dir = tmp_path / "checkpoint" checkpoint_dir.mkdir() bundle_dir = tmp_path / "bundle" - provider = RayTorchRecipeSourceProvider() + provider = RayTorchSourceProvider() with provider.open_source(_checkpoint(checkpoint_dir), _options()) as source: result = BundleExportService().export_bundle( @@ -105,3 +302,28 @@ def test_recipe_source_uses_existing_onnx_bundle_pipeline(tmp_path: Path) -> Non assert result.status == "succeeded" assert (Path(result.canonical_uri) / "manifest.json").is_file() assert any(artifact.format == "onnx" for artifact in result.artifacts) + + +def test_adapter_source_provider_attaches_typed_artifact_plan(tmp_path: Path) -> None: + provider = RayTorchSourceProvider() + options = TorchSourceOptions( + implementation_ref=( + "tests.training.exporters.test_torch_recipe_source:AdapterExportFixture" + ), + implementation_code_digest=CODE_DIGEST, + implementation_id="example.adapter", + policy_digest="1" * 64, + plan_digest="2" * 64, + input_binding_digest="3" * 64, + loop_owner="adapter", + algorithm_config={ + "model": {"input_features": 2}, + "ray": {"storage_path": "/not-visible-to-adapter"}, + "output": {"bundle_uri": "/not-visible-to-adapter"}, + }, + input_bindings={"train": {"feature_names": ["features"]}}, + output_config={"bundle_uri": "/not-visible-to-worker"}, + ) + + with provider.open_source(_adapter_checkpoint(tmp_path), options) as source: + assert source.metadata["artifact_plan"]["roles"] == {"inference": "model"} diff --git a/tests/training/jobs/official_algorithm_gate_job.py b/tests/training/jobs/official_algorithm_gate_job.py index 05c907b..15f4499 100644 --- a/tests/training/jobs/official_algorithm_gate_job.py +++ b/tests/training/jobs/official_algorithm_gate_job.py @@ -70,6 +70,11 @@ def _tensor_columns( f"model input {name!r} has unsupported dynamic trailing shape {shape}" ) width = math.prod(cast(tuple[int, ...], trailing)) if trailing else 1 + if len(trailing) > 1: + # Preserve higher-rank typed tensors as one vector-valued table column; + # flattening them into scalar columns would make inference reconstruct + # rank two and violate the manifest signature. + return (name,) return tuple(f"{name}__{index}" for index in range(width)) @@ -129,15 +134,35 @@ def _stage_bundle_inference_input( columns = _tensor_columns(name=field.name, shape=field.shape) projected.extend(columns) for column_index, column in enumerate(columns): - values[column] = [ - _inference_value( - field.dtype, - row=row, - column=field_index + column_index, - entry_point=entry_point, - ) - for row in range(16) - ] + if len(field.shape[1:]) > 1: + trailing = tuple(cast(tuple[int, ...], field.shape[1:])) + width = math.prod(trailing) + values[column] = [ + np.asarray( + [ + _inference_value( + field.dtype, + row=row, + column=field_index + offset, + entry_point=entry_point, + ) + for offset in range(width) + ] + ) + .reshape(trailing) + .tolist() + for row in range(16) + ] + else: + values[column] = [ + _inference_value( + field.dtype, + row=row, + column=field_index + column_index, + entry_point=entry_point, + ) + for row in range(16) + ] bindings.append( TensorInputBinding( tensor_name=field.name, diff --git a/tests/training/jobs/official_algorithm_matrix.py b/tests/training/jobs/official_algorithm_matrix.py index 71dc73e..2799c42 100644 --- a/tests/training/jobs/official_algorithm_matrix.py +++ b/tests/training/jobs/official_algorithm_matrix.py @@ -29,7 +29,7 @@ class OfficialAlgorithmIdentity: "difference_in_means_ate", "tributo.official.causal.difference_in_means", ), - "dnn.recipe_v2": OfficialAlgorithmIdentity( + "dnn": OfficialAlgorithmIdentity( "tributo-algorithms-tabular-torch", "dnn", "tributo.official.tabular_torch.dnn", @@ -64,10 +64,10 @@ class OfficialAlgorithmIdentity: "graphsage_node_classifier", "tributo.official.graph_pyg.graphsage", ), - "gru_classifier.recipe_v2": OfficialAlgorithmIdentity( + "gru_classifier": OfficialAlgorithmIdentity( "tributo-algorithms-timeseries", "gru_classifier", - "tributo.official.timeseries.gru.recipe_v2", + "tributo.official.timeseries.gru", ), "isolation_forest.parallel_ensemble": OfficialAlgorithmIdentity( "tributo-algorithms-classical", @@ -114,10 +114,10 @@ class OfficialAlgorithmIdentity: "logistic_regression", "tributo.official.logistic_regression.binary_l2", ), - "lstm_classifier.recipe_v2": OfficialAlgorithmIdentity( + "lstm_classifier": OfficialAlgorithmIdentity( "tributo-algorithms-timeseries", "lstm_classifier", - "tributo.official.timeseries.lstm.recipe_v2", + "tributo.official.timeseries.lstm", ), "multinomial_nb": OfficialAlgorithmIdentity( "tributo-algorithms-classical", @@ -139,7 +139,7 @@ class OfficialAlgorithmIdentity: "pretrain_finetune_classifier", "tributo.official.multistage_torch.pretrain_finetune", ), - "pu.recipe_v2": OfficialAlgorithmIdentity( + "pu": OfficialAlgorithmIdentity( "tributo-algorithms-tabular-torch", "pu", "tributo.official.tabular_torch.pu", @@ -252,10 +252,10 @@ def _build_distribution_entry_points() -> Mapping[str, tuple[str, ...]]: "xgboost.framework_native", ), "torch": ( - "dnn.recipe_v2", - "gru_classifier.recipe_v2", - "lstm_classifier.recipe_v2", - "pu.recipe_v2", + "dnn", + "gru_classifier", + "lstm_classifier", + "pu", "tabular_autoencoder", "temporal_conv_classifier", ), diff --git a/tests/training/test_dnn_pu_training.py b/tests/training/test_dnn_pu_training.py index a07d44d..9649b1b 100644 --- a/tests/training/test_dnn_pu_training.py +++ b/tests/training/test_dnn_pu_training.py @@ -304,7 +304,7 @@ def _submit_official_algorithm_gate_job( def test_official_algorithm_wheels_complete_on_ray_cluster( job_client: JobSubmissionClient, ) -> None: - """Prove official decomposition and RecipeV2 Wheels train across nodes.""" + """Prove official decomposition and Torch Wheels train across nodes.""" if os.environ.get("TRIBUTO_DOCKER_DISTRIBUTED_ALGORITHM_IT") != "1": pytest.fail("official algorithm IT must run in its owned Docker cluster") wheels = _official_algorithm_wheels() @@ -419,7 +419,7 @@ def test_official_algorithm_wheels_complete_on_ray_cluster( "ray_parallel_ensemble", "ray_iterative_optimization", "ray_map_reduce", - "ray_train_recipe_v2", + "ray_train_torch", "framework_native", } actual_strategies = {record["receipt"]["strategy"] for record in records} diff --git a/tests/training/test_training_lifecycle.py b/tests/training/test_training_lifecycle.py index 6c43d5a..0bf60cf 100644 --- a/tests/training/test_training_lifecycle.py +++ b/tests/training/test_training_lifecycle.py @@ -714,7 +714,7 @@ def test_first_party_source_providers_do_not_require_entry_points( lifecycle_module._load_provider_plugins(registry) - assert registry.list_all() == ["ray-torch-recipe-v1"] + assert registry.list_all() == ["ray-torch-v1"] def test_bundle_real_routing_contract( self, monkeypatch: pytest.MonkeyPatch diff --git a/tools/check_public_api_annotations.py b/tools/check_public_api_annotations.py index b736ba1..7b0618a 100644 --- a/tools/check_public_api_annotations.py +++ b/tools/check_public_api_annotations.py @@ -46,6 +46,9 @@ # their concrete handle types are annotated. "DataHandle", "WriteHandle", + # Torch loss union aliases — type aliases cannot be decorated; their + # concrete loss contribution types are annotated. + "TorchStepLoss", # Stable string constants — Python str objects cannot be decorated. "ARTIFACT_KIND_MODEL", "ARTIFACT_KIND_REPORT", @@ -62,6 +65,9 @@ "TORCH", "TRANSFORMERS", "XGBOOST", + # Runtime identity constants are stable strings and cannot carry a + # decorator; the Runtime class and entrypoint functions are annotated. + "RAY_TRAIN_TORCH_RUNTIME_ID", } From cdb247c801dee816fdcccdd96d1ddc6f59f7d77d Mon Sep 17 00:00:00 2001 From: jiangxt2 Date: Fri, 4 Sep 2026 23:15:08 +0800 Subject: [PATCH 2/2] refactor(algorithms): simplify PyTorch runtime contracts Remove custom recovery and storage control paths, keep checkpoint lifecycle on Ray, and tighten Torch evidence, routing, and installed-distribution conformance. Signed-off-by: jiangxt2 --- docs/architecture/ray-first-torch-recipes.md | 69 +- docs/examples/doc_code/pu_execution.json | 3 +- docs/how-to/custom-distributed-algorithms.md | 60 +- docs/reference/api/algorithms-training.md | 48 - src/tributo/algorithms/__init__.py | 24 - src/tributo/algorithms/api/__init__.py | 24 - src/tributo/algorithms/api/distribution.py | 290 +-- src/tributo/algorithms/api/execution.py | 28 +- src/tributo/algorithms/api/models.py | 22 - src/tributo/algorithms/api/support.py | 5 +- src/tributo/algorithms/api/torch_runtime.py | 1044 +------- src/tributo/algorithms/conformance.py | 102 + src/tributo/algorithms/conformance_cli.py | 48 + src/tributo/algorithms/core/builder.py | 15 +- src/tributo/algorithms/core/dispatcher.py | 138 +- src/tributo/algorithms/core/planner.py | 5 - src/tributo/algorithms/spi/__init__.py | 2 - src/tributo/algorithms/spi/execution.py | 18 - src/tributo/algorithms/spi/torch.py | 71 +- src/tributo/inference/kernel.py | 22 - .../algorithm_runtimes/ray_train_torch.py | 2107 +++-------------- src/tributo/integrations/sources/ray_torch.py | 33 +- src/tributo/training/portable_tune.py | 12 + tests/algorithms/test_conformance_cli.py | 109 + tests/algorithms/test_execution_contract.py | 18 + tests/algorithms/test_support_evidence.py | 30 + tests/algorithms/test_torch_recipe_worker.py | 272 ++- .../test_torch_runtime_contract_v1.py | 496 ++-- tests/inference/test_contracts.py | 50 - .../test_distributed_algorithm_it_contract.py | 43 + .../exporters/test_torch_recipe_source.py | 29 +- .../jobs/official_algorithm_gate_job.py | 51 +- .../jobs/official_algorithm_identities.json | 42 + .../jobs/official_algorithm_matrix.py | 233 +- tests/training/test_portable_tune.py | 11 + tools/algorithm_gate_provenance.py | 30 + 36 files changed, 1496 insertions(+), 4108 deletions(-) create mode 100644 src/tributo/algorithms/conformance_cli.py create mode 100644 tests/algorithms/test_conformance_cli.py create mode 100644 tests/training/jobs/official_algorithm_identities.json diff --git a/docs/architecture/ray-first-torch-recipes.md b/docs/architecture/ray-first-torch-recipes.md index cdedf74..ee04317 100644 --- a/docs/architecture/ray-first-torch-recipes.md +++ b/docs/architecture/ray-first-torch-recipes.md @@ -20,8 +20,8 @@ implementation and its public annotations, not unversioned latest docs. | Torch batches | Public `DataIterator.iter_torch_batches()` | Validate feature/label roles and pass bounded batch options | | Device and DDP | Stable `ray.train.torch.prepare_model()` | Reject unvalidated BatchNorm and GPU claims | | Metrics | PyTorch tensors and `torch.distributed` collectives | Apply the descriptor's existing `MetricReduction`; do not infer reducers from names | -| Checkpoint transfer | Ray `Checkpoint`, `train.report()`, `RunConfig`, and retention | Write bounded model, optimizer, scheduler, gradient scaling, and RNG metadata at complete optimizer-window boundaries and validate Stage identity | -| Run identity | V2 `RunConfig.name` and storage context | Derive a unique Ray run name from `run_id`, `invocation_id`, `stage_id`, and the Policy/execution-plan identity digest so retries reuse only the matching Stage directory | +| Checkpoint transfer | Ray `Checkpoint`, `train.report()`, `RunConfig`, and retention | Report one completed-Stage checkpoint for Stage dependencies and export; retry replays the Stage | +| Run identity | V2 `RunConfig.name` and storage context | Derive a unique Ray run name from the logical run, Stage, and Policy/execution-plan identity | | Export and serving | Existing Torch ONNX exporter, ONNX Runtime validator, Bundle publisher, reader, batch inference, and Serve flavor | Reconstruct the trusted recipe model and create one `ExportSource` | | Local execution | `ray.init(address="local")` | Own and close only the runtime created by Tributo | | Existing or provisioned cluster | `ray.init(address="auto")`, Ray Jobs, KubeRay RayJob, or Cluster Launcher | Attach or expose the same entrypoint; workload code never provisions or deletes the cluster | @@ -33,18 +33,16 @@ changes only the `equal` argument while preserving Ray execution options, training-resource exclusion, and locality hints. A future Ray public option should replace this adapter. -Ray Train V2 rejects a non-null `resume_from_checkpoint` argument as deprecated. -The Core Runtime never supplies it: failure retry uses `train.get_checkpoint()`; -cross-Stage and cross-Run recovery use a credential-free Core Worker control -envelope. A full `TorchRecoveryEnvelope` records completed Stage locators and an -optional active Stage; completed Stage evidence is persisted in the validated -checkpoint sidecar so a recovery-only invocation can produce the same Receipt. -A typed progress manifest preserves per-rank coverage/metric prefixes, reducer -observation, Dataset cursor and whether the current epoch's Scheduler step has -already run. Remote Stage payloads remain unavailable under a deterministic -staging prefix until a matching commit marker is written. -A Torch-only preflight and invocation-local one-shot lease complete before Ray -or input resources are opened. +Ray Train V2 rejects a non-null `resume_from_checkpoint` argument as deprecated, +so the Core Runtime never supplies it. Ray owns Worker failure retry and +Checkpoint persistence. Runtime API v1 replays the current Recipe or Adapter +Stage from its Dataset beginning instead of exposing Ray's retry Checkpoint to +algorithm code; a replayed component Stage still receives its predecessor Stage +Checkpoint. Core does not add a progress cursor or remote commit protocol. +Completed Stage checkpoints move directly through Ray objects to dependent +Stages and the Bundle exporter. Cross-Run recovery is not supported. Torch-only +environment validation still runs before Ray or input resources are opened, +without a separate token or lease state machine. ## Uneven input decision @@ -56,16 +54,15 @@ gradient. Join does not mirror that additional collective. The first implementation therefore uses a narrow dynamic alignment protocol: - every rank must receive at least one batch; -- all ranks exchange only the active row count before each step; +- all ranks exchange only an active-batch flag before each step; - exhausted ranks run a zero-row DDP forward/backward and contribute zero gradient while active ranks retain every observed row; - loss scaling uses the algorithm-declared global normalizer for the complete accumulation window, so uneven final batches remain mathematically weighted by the declared unit (rows, sample weights, or valid tokens); -- no row is dropped or replayed; -- all ranks report at the declared `checkpoint_interval_windows` cadence (one - complete optimizer window by default), and only the checkpoint-owner rank - attaches the replicated checkpoint. +- no row is dropped or replayed within a successful Stage attempt; +- all ranks report once when the Stage completes, and only the checkpoint-owner + rank attaches the replicated checkpoint. Optional validation and test roles use the same rank-alignment rule. When a present role has fewer rows than Workers, exhausted ranks execute a typed @@ -84,11 +81,11 @@ test datasets are optional roles when the input adapter can provide them; the recipe does not split data. The first-party DNN recipe keeps its existing split behavior in a thin adapter before the common worker loop. -The first `tabular_batch` profile accepts scalar numeric feature columns. The -worker stacks them into one dense tensor, while the export adapter restores the -original column names as separate ONNX inputs so Bundle inference does not lose -the declared `InputBinding` contract. Vector, sequence, jagged, and graph inputs -remain later profiles rather than implicit shape inference. +Recipe implementations construct their declared dense, fixed-window, token or +ID tensors from named scalar columns. LSTM and GRU export a two-dimensional +`[batch, window]` input and add the singleton channel inside the model, keeping +the existing format-neutral inference binding unchanged. Jagged and graph +shapes remain explicit Adapter contracts rather than implicit Core inference. Advanced models implement the typed `TorchRecipe` hooks. Models that need custom training steps, framework callbacks, graph sampling, sharded embeddings, @@ -99,21 +96,21 @@ An Adapter declares its `TorchArtifactPlan` through the Core Provider. The Provider attaches that plan to the Adapter's `ExportSource` before invoking the single Core Bundle exporter, and rejects conflicting plan metadata. Adapter worker configuration receives algorithm-owned values plus credential-free input -binding metadata; Core Ray paths, recovery locators, and output publication -URIs are never accepted in that JSON payload. +binding metadata; Core Ray paths and output publication URIs are never accepted +in that JSON payload. Capabilities not described by this contract are not implied by the recipe, dependency set, or framework installation. ## Evidence status -Unit evidence covers descriptor lowering, exact-coverage configuration, -streaming batch consumption, typed metric reduction, resumable checkpoint -contents, trusted model reconstruction, ONNX Bundle publication, execution -profile compatibility, and isolated wheel construction. The isolated Recipe -wheel has also passed owned-local single/two-worker and existing Docker Ray -cluster multi-node Ray Jobs gates, including 65-row uneven input and Bundle -publication. Kubernetes and KubeRay remain external deployment substrates; the -deployment-neutral Docker Ray multi-node gate is the Tributo evidence for the -common `cluster` workload. This does not claim that Tributo owns or validates -Kubernetes control-plane lifecycle. +Local contract evidence covers descriptor lowering, exact-coverage +configuration, streaming batch consumption, typed metric reduction, completed +Stage checkpoints, trusted model reconstruction, Bundle publication, and +execution-profile compatibility. Candidate-Wheel conformance and the final +Docker Ray multi-node gate remain required before the migrated algorithms are +declared validated. Kubernetes and KubeRay remain external deployment +substrates; Tributo does not own or validate their control-plane lifecycle. +Portable Torch Tune trials report the target metric only; the selected +parameters must be used for a separate formal training run that publishes the +final Bundle. diff --git a/docs/examples/doc_code/pu_execution.json b/docs/examples/doc_code/pu_execution.json index 9a999a6..c5227de 100644 --- a/docs/examples/doc_code/pu_execution.json +++ b/docs/examples/doc_code/pu_execution.json @@ -34,8 +34,7 @@ }, "ray": { "max_failures": 0, - "storage_path": "/shared/ray-results/pu-fraud", - "resume": {"checkpoint_interval": 1} + "storage_path": "/shared/ray-results/pu-fraud" }, "output": { "bundle_uri": "/shared/models/pu-fraud" diff --git a/docs/how-to/custom-distributed-algorithms.md b/docs/how-to/custom-distributed-algorithms.md index 969f666..b3278cc 100644 --- a/docs/how-to/custom-distributed-algorithms.md +++ b/docs/how-to/custom-distributed-algorithms.md @@ -215,6 +215,40 @@ artifact preflight. A non-empty `algorithm_ids` allowlist is enforced against an offline Bundle's manifest; an empty allowlist means that the Profile does not restrict the algorithm family. +## Migrate legacy PyTorch extensions + +The unified Torch API replaces the former recipe and framework-native PyTorch +entry points. The migration is intentionally source-breaking; there are no +legacy aliases or compatibility adapters. + +| Legacy API or hook | Unified Torch replacement | +| --- | --- | +| `TorchTrainingRecipe.model_factory()` and `loss_factory()` | `TorchRecipe.build_modules(TorchBuildContext)` returning a `TorchModuleSet` | +| `TorchTrainingRecipe.optimizer_factory()` | `TorchRecipe.configure_optimizers()` returning `TorchOptimizationPlan` | +| `TorchTrainingRecipe.metric_factories()` | Typed contributions from `training_step()` / `validation_step()` plus declarations from `metric_plan()` | +| `TorchTrainingRecipe.forward()` and `compute_loss()` | Algorithm-owned `training_step()` / `validation_step()` | +| `TrainingRecipeV2.build_modules(config)` | `TorchRecipe.build_modules(TorchBuildContext)` | +| `TrainingRecipeV2.batch_adapter()` | `TorchRecipe.adapt_batch(batch, TorchBatchContext)` returning `TorchBatch` | +| `TrainingRecipeV2.training_step()` and `validation_step()` | The same named typed hooks returning `TorchStepResult` with explicit numerator/normalizer contributions | +| `TrainingRecipeV2.optimization_plan()` | `TorchRecipe.configure_optimizers()` | +| `TrainingRecipeV2.metric_plan()` callable factories | `TorchRecipe.metric_plan()` reducer declarations; step hooks return metric contributions | +| `TrainingRecipeV2.checkpoint_codec()` | Core-owned checkpoint reporting and `TorchCheckpointDescriptor`; no Recipe-owned serialization hook | +| PyTorch `FrameworkNativeAlgorithm.validate_environment()` and `bind_datasets()` | `RayTorchAdapter.validate_environment(context)` and `bind_datasets(datasets, context)` | +| PyTorch `FrameworkNativeAlgorithm.build_trainer()` | Removed; Core owns the sole `TorchTrainer`, while the Adapter supplies `worker_config()` and `train_loop_per_worker()` | +| PyTorch `FrameworkNativeAlgorithm.collect_evidence()` | Core-owned evidence collection from bounded worker reports | +| PyTorch `FrameworkNativeAlgorithm.checkpoint_source()` | `RayTorchAdapter.checkpoint_source(result, context)` | +| `RayTorchRecipeSourceProvider` / `TorchRecipeSourceOptions` | `RayTorchSourceProvider` / `TorchSourceOptions`; Recipes use Core reconstruction, while Adapters provide `artifact_plan()` and `open_export_source()` | + +Registration and runtime identities migrate as follows: + +| Legacy registration | Unified registration | +| --- | --- | +| `AlgorithmBuilder.from_torch_recipe()` | `AlgorithmBuilder.from_torch()` | +| `AlgorithmBuilder.from_training_recipe_v2()` | `AlgorithmBuilder.from_torch()` | +| PyTorch `FrameworkNativeAlgorithm` registrations | `AlgorithmBuilder.from_torch_adapter()` with `RAY_TRAIN_TORCH` | +| Legacy PyTorch Runtime IDs | `tributo.ray_train_torch` | +| `ray-torch-recipe-v1` Source Provider | `ray-torch-v1` | + ## Use a low-code PyTorch recipe For Core-owned training, subclass `TorchRecipe` and implement the typed @@ -226,9 +260,11 @@ and cannot create a nested Trainer or declare a second execution plan. ```python from tributo.algorithms import ( + TorchArtifactPlan, TorchBatch, TorchLossContribution, TorchMetricPlan, + TorchModuleSet, TorchOptimizationPlan, TorchRecipe, TorchStepResult, @@ -239,13 +275,15 @@ class BinaryLinearRecipe(TorchRecipe): def build_modules(self, context): import torch - return {"model": torch.nn.Linear(2, 1), "loss": torch.nn.BCEWithLogitsLoss()} + return TorchModuleSet( + {"model": torch.nn.Linear(2, 1), "loss": torch.nn.BCEWithLogitsLoss()} + ) def adapt_batch(self, batch, context): import torch features = torch.as_tensor(batch["features"]) - targets = torch.as_tensor(batch["label"]) + targets = torch.as_tensor(batch["label"]).reshape(-1, 1) return TorchBatch(positional=(features,), targets=targets, local_rows=len(targets)) def training_step(self, modules, batch, context): @@ -272,7 +310,23 @@ class BinaryLinearRecipe(TorchRecipe): return TorchMetricPlan({"train_loss": "sum_count"}) def artifact_plan(self, context): - return {"source_kind": "torch_module", "roles": {"inference": "onnx-model"}} + return TorchArtifactPlan( + source_kind="torch_module", + input_signature=( + {"name": "features", "dtype": "float32", "shape": ("batch", 2)}, + ), + output_signature=( + {"name": "prediction", "dtype": "float32", "shape": ("batch", 1)}, + ), + targets=( + { + "name": "onnx-model", + "format": "onnx", + "exporter_id": "torch-onnx-v1", + }, + ), + roles={"inference": "onnx-model"}, + ) ``` Use `AlgorithmBuilder.from_torch()` to lower a Recipe, or diff --git a/docs/reference/api/algorithms-training.md b/docs/reference/api/algorithms-training.md index 8779c2a..c547350 100644 --- a/docs/reference/api/algorithms-training.md +++ b/docs/reference/api/algorithms-training.md @@ -334,18 +334,10 @@ documentation for every public stability tier. :no-members: ``` -```{autoclass} tributo.algorithms.api.torch_runtime.TorchCheckpointLocator -:no-members: -``` - ```{autoclass} tributo.algorithms.api.torch_runtime.TorchCheckpointPayloadDraft :no-members: ``` -```{autoclass} tributo.algorithms.api.torch_runtime.TorchCheckpointProgress -:no-members: -``` - ```{autoclass} tributo.algorithms.api.torch_runtime.TorchCheckpointRef :no-members: ``` @@ -390,43 +382,13 @@ documentation for every public stability tier. :no-members: ``` -```{autoclass} tributo.algorithms.api.torch_runtime.TorchPreflightLease -:no-members: -``` - -```{autoclass} tributo.algorithms.api.torch_runtime.TorchPreflightTokenData -:no-members: -``` - -```{autoclass} tributo.algorithms.api.torch_runtime.TorchRankProgressStatistics -:no-members: -``` - -```{autoclass} tributo.algorithms.api.torch_runtime.TorchRecoveryEnvelope -:no-members: -``` - -```{autoclass} tributo.algorithms.api.torch_runtime.TorchRuntimeExecutionEnvelope -:no-members: -``` - ```{autoclass} tributo.algorithms.api.torch_runtime.TorchStageRunIdentity :no-members: ``` -```{autoclass} tributo.algorithms.api.torch_runtime.TorchWorkerControlEnvelope -:no-members: -``` - ```{autofunction} tributo.algorithms.api.torch_runtime.apply_torch_loss_backward ``` -```{autofunction} tributo.algorithms.api.torch_runtime.claim_torch_run_directory -``` - -```{autofunction} tributo.algorithms.api.torch_runtime.describe_torch_checkpoint -``` - ```{autofunction} tributo.algorithms.api.torch_runtime.invoke_torch_global_loss_reducer ``` @@ -436,12 +398,6 @@ documentation for every public stability tier. ```{autofunction} tributo.algorithms.api.torch_runtime.report_torch_checkpoint ``` -```{autofunction} tributo.algorithms.api.torch_runtime.torch_run_config_name -``` - -```{autofunction} tributo.algorithms.api.torch_runtime.validate_torch_retry_identity -``` - ## `tributo.algorithms.composition` @@ -545,10 +501,6 @@ documentation for every public stability tier. :no-members: ``` -```{autoclass} tributo.algorithms.spi.execution.TorchRuntimePreflight -:no-members: -``` - ```{autoclass} tributo.algorithms.spi.execution.Transformable :no-members: ``` diff --git a/src/tributo/algorithms/__init__.py b/src/tributo/algorithms/__init__.py index 6549028..84cca9d 100644 --- a/src/tributo/algorithms/__init__.py +++ b/src/tributo/algorithms/__init__.py @@ -56,9 +56,7 @@ TorchBackwardContext, TorchBackwardResult, TorchCheckpointDescriptor, - TorchCheckpointLocator, TorchCheckpointPayloadDraft, - TorchCheckpointProgress, TorchCheckpointRef, TorchCompositeGlobalState, TorchCompositeLossContribution, @@ -74,28 +72,18 @@ TorchMetricReductionContext, TorchMetricReductionResult, TorchPolicy, - TorchPreflightLease, - TorchPreflightTokenData, - TorchRankProgressStatistics, - TorchRecoveryEnvelope, TorchRoleExecutionEvidence, - TorchRuntimeExecutionEnvelope, TorchStageRunIdentity, TorchStageSpec, TorchStepLoss, - TorchWorkerControlEnvelope, UserExecutionContext, WorkerExecutionEvidence, WorkerRange, WorkerResources, apply_torch_loss_backward, - claim_torch_run_directory, - describe_torch_checkpoint, invoke_torch_global_loss_reducer, reduce_torch_metrics, report_torch_checkpoint, - torch_run_config_name, - validate_torch_retry_identity, ) from tributo.algorithms.composition import build_algorithm_dispatcher from tributo.algorithms.conformance import ( @@ -215,9 +203,7 @@ "TorchBackwardContext", "TorchBackwardResult", "TorchCheckpointPayloadDraft", - "TorchCheckpointProgress", "TorchCheckpointDescriptor", - "TorchCheckpointLocator", "TorchCheckpointRef", "TorchCompositeGlobalState", "TorchCompositeLossContribution", @@ -229,22 +215,12 @@ "TorchMetricPolicy", "TorchMetricReductionContext", "TorchMetricReductionResult", - "TorchPreflightLease", - "TorchPreflightTokenData", - "TorchRecoveryEnvelope", - "TorchRankProgressStatistics", - "TorchRuntimeExecutionEnvelope", "TorchStageRunIdentity", "TorchStepLoss", - "TorchWorkerControlEnvelope", "apply_torch_loss_backward", - "claim_torch_run_directory", - "describe_torch_checkpoint", "invoke_torch_global_loss_reducer", "reduce_torch_metrics", "report_torch_checkpoint", - "torch_run_config_name", - "validate_torch_retry_identity", "build_algorithm_dispatcher", "validate_algorithm_descriptor_conformance", "validate_installed_algorithm_package", diff --git a/src/tributo/algorithms/api/__init__.py b/src/tributo/algorithms/api/__init__.py index 873a272..9005eb7 100644 --- a/src/tributo/algorithms/api/__init__.py +++ b/src/tributo/algorithms/api/__init__.py @@ -90,9 +90,7 @@ TorchBackwardContext, TorchBackwardResult, TorchCheckpointDescriptor, - TorchCheckpointLocator, TorchCheckpointPayloadDraft, - TorchCheckpointProgress, TorchCheckpointRef, TorchCompositeGlobalState, TorchCompositeLossContribution, @@ -104,22 +102,12 @@ TorchMetricPolicy, TorchMetricReductionContext, TorchMetricReductionResult, - TorchPreflightLease, - TorchPreflightTokenData, - TorchRankProgressStatistics, - TorchRecoveryEnvelope, - TorchRuntimeExecutionEnvelope, TorchStageRunIdentity, TorchStepLoss, - TorchWorkerControlEnvelope, apply_torch_loss_backward, - claim_torch_run_directory, - describe_torch_checkpoint, invoke_torch_global_loss_reducer, reduce_torch_metrics, report_torch_checkpoint, - torch_run_config_name, - validate_torch_retry_identity, ) __all__ = [ @@ -200,9 +188,7 @@ "TorchBackwardContext", "TorchBackwardResult", "TorchCheckpointPayloadDraft", - "TorchCheckpointProgress", "TorchCheckpointDescriptor", - "TorchCheckpointLocator", "TorchCheckpointRef", "TorchCompositeGlobalState", "TorchCompositeLossContribution", @@ -214,20 +200,10 @@ "TorchMetricPolicy", "TorchMetricReductionContext", "TorchMetricReductionResult", - "TorchPreflightLease", - "TorchPreflightTokenData", - "TorchRecoveryEnvelope", - "TorchRankProgressStatistics", - "TorchRuntimeExecutionEnvelope", "TorchStageRunIdentity", "TorchStepLoss", - "TorchWorkerControlEnvelope", "apply_torch_loss_backward", - "claim_torch_run_directory", - "describe_torch_checkpoint", "invoke_torch_global_loss_reducer", "reduce_torch_metrics", "report_torch_checkpoint", - "torch_run_config_name", - "validate_torch_retry_identity", ] diff --git a/src/tributo/algorithms/api/distribution.py b/src/tributo/algorithms/api/distribution.py index 5e81792..4364b6b 100644 --- a/src/tributo/algorithms/api/distribution.py +++ b/src/tributo/algorithms/api/distribution.py @@ -606,66 +606,27 @@ class TorchStageSpec: """One Core-orchestrated stage in a Torch execution plan.""" stage_id: str - worker_loop_ref: str input_roles: tuple[str, ...] - depends_on: tuple[str, ...] = () checkpoint_from_stage: str | None = None - metric_mapping: Mapping[str, str] = field(default_factory=dict) - checkpoint_required: bool = True - checkpoint_interval_windows: int = 1 def __post_init__(self) -> None: _string(self.stage_id, "Torch stage_id") - _qualified_reference(self.worker_loop_ref, "Torch worker_loop_ref") roles = tuple(self.input_roles) if not roles or any(not isinstance(role, str) or not role for role in roles): raise AlgorithmConfigurationError("Torch stage input_roles are required") if len(set(roles)) != len(roles): raise AlgorithmConfigurationError("Torch stage input_roles must be unique") - depends = tuple(self.depends_on) - if any(not isinstance(item, str) or not item for item in depends): - raise AlgorithmConfigurationError("Torch stage dependencies are invalid") - if len(set(depends)) != len(depends): - raise AlgorithmConfigurationError("Torch stage dependencies must be unique") if self.checkpoint_from_stage is not None and not isinstance( self.checkpoint_from_stage, str ): raise AlgorithmConfigurationError("Torch checkpoint_from_stage is invalid") - _boolean(self.checkpoint_required, "Torch checkpoint_required") - _positive_integer( - self.checkpoint_interval_windows, - "Torch checkpoint_interval_windows", - ) - if any( - not isinstance(name, str) - or not name - or not isinstance(value, str) - or not value - for name, value in self.metric_mapping.items() - ): - raise AlgorithmConfigurationError( - "Torch stage metric_mapping must be named strings" - ) - if len(set(self.metric_mapping.values())) != len(self.metric_mapping): - raise AlgorithmConfigurationError( - "Torch stage metric_mapping targets must be unique" - ) object.__setattr__(self, "input_roles", roles) - object.__setattr__(self, "depends_on", depends) - object.__setattr__( - self, "metric_mapping", FrozenDict(dict(self.metric_mapping)) - ) def to_dict(self) -> dict[str, Any]: payload: dict[str, Any] = { "stage_id": self.stage_id, - "worker_loop_ref": self.worker_loop_ref, "input_roles": list(self.input_roles), - "depends_on": list(self.depends_on), "checkpoint_from_stage": self.checkpoint_from_stage, - "metric_mapping": dict(self.metric_mapping), - "checkpoint_required": self.checkpoint_required, - "checkpoint_interval_windows": self.checkpoint_interval_windows, } return payload @@ -708,8 +669,6 @@ class SingleStageTorchPlan(TorchExecutionPlan): stage: TorchStageSpec = field( default_factory=lambda: TorchStageSpec( "train", - "tributo.integrations.algorithm_runtimes.ray_train_torch:" - "torch_recipe_train_loop_per_worker", ("train",), ) ) @@ -719,7 +678,7 @@ def __post_init__(self) -> None: raise AlgorithmConfigurationError( "Torch execution plan api_version must be 1" ) - if self.stage.depends_on or self.stage.checkpoint_from_stage is not None: + if self.stage.checkpoint_from_stage is not None: raise AlgorithmConfigurationError( "single Torch stage cannot have dependencies" ) @@ -751,10 +710,6 @@ def __post_init__(self) -> None: raise AlgorithmConfigurationError("component Torch stage IDs are invalid") prior: set[str] = set() for stage in stages: - if any(dep not in prior for dep in stage.depends_on): - raise AlgorithmConfigurationError( - "Torch stage dependencies must reference earlier stages" - ) if ( stage.checkpoint_from_stage is not None and stage.checkpoint_from_stage not in prior @@ -762,16 +717,6 @@ def __post_init__(self) -> None: raise AlgorithmConfigurationError( "Torch checkpoint_from_stage must reference an earlier stage" ) - if stage.checkpoint_from_stage is not None: - source = next( - item - for item in stages - if item.stage_id == stage.checkpoint_from_stage - ) - if not source.checkpoint_required: - raise AlgorithmConfigurationError( - "Torch checkpoint_from_stage requires a checkpoint-producing source" - ) prior.add(stage.stage_id) object.__setattr__(self, "stages", stages) @@ -798,9 +743,7 @@ class TorchPolicy: metric_reducers: Mapping[str, MetricReduction] backend: str = "auto" checkpoint_owner_rank: int = 0 - resume_supported: bool = True - same_world_size_resume: bool | None = True - rank_seeded: bool = True + resume_supported: bool = False checkpoint_adapter_ref: str | None = None evidence_adapter_ref: str | None = None global_loss_reducer_ref: str | None = None @@ -884,16 +827,9 @@ def __post_init__(self) -> None: raise AlgorithmConfigurationError("Torch backend is invalid") _non_negative_integer(self.checkpoint_owner_rank, "Torch checkpoint_owner_rank") _boolean(self.resume_supported, "Torch resume_supported") - _boolean(self.rank_seeded, "Torch rank_seeded") - if self.same_world_size_resume is not None: - _boolean(self.same_world_size_resume, "Torch same_world_size_resume") - if self.resume_supported and self.same_world_size_resume is not True: - raise AlgorithmConfigurationError( - "Torch v1 recovery only supports same-world-size resume" - ) - if not self.resume_supported and self.same_world_size_resume is not None: + if self.resume_supported: raise AlgorithmConfigurationError( - "unsupported Torch resume must omit same_world_size_resume" + "Torch Runtime API v1 does not support cross-Run recovery" ) if self.max_replicated_bytes_per_worker is not None: _positive_integer( @@ -966,7 +902,6 @@ def to_dict(self) -> dict[str, Any]: "backend": self.backend, "checkpoint_owner_rank": self.checkpoint_owner_rank, "resume_supported": self.resume_supported, - "rank_seeded": self.rank_seeded, "checkpoint_adapter_ref": self.checkpoint_adapter_ref, "evidence_adapter_ref": self.evidence_adapter_ref, "global_loss_reducer_ref": self.global_loss_reducer_ref, @@ -976,8 +911,6 @@ def to_dict(self) -> dict[str, Any]: "capabilities": list(self.capabilities), "max_replicated_bytes_per_worker": self.max_replicated_bytes_per_worker, } - if self.same_world_size_resume is not None: - payload["same_world_size_resume"] = self.same_world_size_resume return payload @property @@ -989,43 +922,46 @@ def digest(self) -> str: def from_dict(cls, value: Mapping[str, Any]) -> "TorchPolicy": """Reconstruct a policy without importing an implementation module.""" try: + if set(value) & {"same_world_size_resume", "rank_seeded"}: + raise AlgorithmConfigurationError( + "TorchPolicy payload contains removed recovery fields" + ) plan_value = _mapping(value["execution_plan"], "Torch execution_plan") raw_stages = _sequence(plan_value["stages"], "Torch execution stages") + if any( + set(_mapping(item, "Torch stage")) + & { + "worker_loop_ref", + "depends_on", + "metric_mapping", + "checkpoint_required", + } + for item in raw_stages + ): + raise AlgorithmConfigurationError( + "Torch stage payload contains removed fields" + ) stages = tuple( TorchStageSpec( stage_id=_string( _mapping(item, "Torch stage")["stage_id"], "stage_id" ), - worker_loop_ref=_string( - _mapping(item, "Torch stage")["worker_loop_ref"], - "worker_loop_ref", - ), input_roles=tuple( _sequence( _mapping(item, "Torch stage")["input_roles"], "input_roles" ) ), - depends_on=tuple( - _sequence( - _mapping(item, "Torch stage").get("depends_on", ()), - "depends_on", - ) - ), checkpoint_from_stage=_mapping(item, "Torch stage").get( "checkpoint_from_stage" ), - metric_mapping=_mapping( - _mapping(item, "Torch stage").get("metric_mapping", {}), - "metric_mapping", - ), - checkpoint_required=_boolean( - _mapping(item, "Torch stage").get("checkpoint_required", True), - "checkpoint_required", - ), ) for item in raw_stages ) if plan_value.get("kind") == "single": + if len(stages) != 1: + raise AlgorithmConfigurationError( + "single Torch execution plan requires one stage" + ) execution_plan: TorchExecutionPlan = SingleStageTorchPlan( api_version=plan_value["api_version"], stage=stages[0] ) @@ -1059,9 +995,7 @@ def from_dict(cls, value: Mapping[str, Any]) -> "TorchPolicy": metric_reducers=metric_reducers, backend=value.get("backend", "auto"), checkpoint_owner_rank=value.get("checkpoint_owner_rank", 0), - resume_supported=value.get("resume_supported", True), - same_world_size_resume=value.get("same_world_size_resume"), - rank_seeded=value.get("rank_seeded", True), + resume_supported=value.get("resume_supported", False), checkpoint_adapter_ref=value.get("checkpoint_adapter_ref"), evidence_adapter_ref=value.get("evidence_adapter_ref"), global_loss_reducer_ref=value.get("global_loss_reducer_ref"), @@ -1447,172 +1381,12 @@ def from_dict(cls, value: Mapping[str, Any]) -> DistributionSpec: exactness=DistributedExactness(policy_value["exactness"]), ) elif kind == "torch": - execution_plan = _mapping( - policy_value["execution_plan"], "Torch execution_plan" - ) - stage_values = _sequence( - execution_plan["stages"], "Torch execution stages" - ) - stages = tuple( - TorchStageSpec( - stage_id=_string( - _mapping(item, "Torch stage")["stage_id"], - "Torch stage_id", - ), - worker_loop_ref=_string( - _mapping(item, "Torch stage")["worker_loop_ref"], - "Torch worker_loop_ref", - ), - input_roles=tuple( - _string(role, "Torch input role") - for role in _sequence( - _mapping(item, "Torch stage")["input_roles"], - "Torch input_roles", - ) - ), - depends_on=tuple( - _string(dep, "Torch dependency") - for dep in _sequence( - _mapping(item, "Torch stage").get("depends_on", ()), - "Torch depends_on", - ) - ), - checkpoint_from_stage=_mapping(item, "Torch stage").get( - "checkpoint_from_stage" - ), - metric_mapping=_mapping( - _mapping(item, "Torch stage").get("metric_mapping", {}), - "Torch metric_mapping", - ), - checkpoint_required=_boolean( - _mapping(item, "Torch stage").get( - "checkpoint_required", True - ), - "Torch checkpoint_required", - ), - ) - for item in stage_values - ) - plan_value: TorchExecutionPlan - if execution_plan.get("kind") == "single": - plan_value = SingleStageTorchPlan( - api_version=_positive_integer( - execution_plan["api_version"], - "Torch execution plan api_version", - ), - stage=stages[0], - ) - else: - plan_value = ComponentStageTorchPlan( - api_version=_positive_integer( - execution_plan["api_version"], - "Torch execution plan api_version", - ), - stages=stages, - final_stage_id=_string( - execution_plan["final_stage_id"], - "Torch final_stage_id", - ), - ) - routes = tuple( - TorchDatasetRoute( - role=_string( - _mapping(item, "Torch route")["role"], "Torch role" - ), - mode=_string( - _mapping(item, "Torch route")["mode"], "Torch mode" - ), - required=_boolean( - _mapping(item, "Torch route").get("required", True), - "Torch route required", - ), - min_total_rows_if_present=_non_negative_integer( - _mapping(item, "Torch route").get( - "min_total_rows_if_present", 1 - ), - "Torch minimum total rows", - ), - min_rows_per_worker=_non_negative_integer( - _mapping(item, "Torch route").get("min_rows_per_worker", 1), - "Torch minimum rows per worker", - ), - empty_rank_policy=_string( - _mapping(item, "Torch route").get( - "empty_rank_policy", "reject" - ), - "Torch empty rank policy", - ), - max_rows=_mapping(item, "Torch route").get("max_rows"), - max_bytes_per_worker=_mapping(item, "Torch route").get( - "max_bytes_per_worker" - ), - ) - for item in _sequence( - policy_value["dataset_routing"], "Torch dataset_routing" - ) - ) - policy = TorchPolicy( - torch_runtime_api_version=_positive_integer( - policy_value["torch_runtime_api_version"], - "Torch Runtime API version", - ), - loop_owner=_string(policy_value["loop_owner"], "Torch loop_owner"), - parallelism_id=_string( - policy_value["parallelism_id"], "Torch parallelism_id" - ), - dataset_routing=routes, - execution_plan=plan_value, - state_layout=_string( - policy_value["state_layout"], "Torch state_layout" - ), - metric_reducers={ - _string(name, "Torch metric name"): MetricReduction(value) - for name, value in _mapping( - policy_value["metric_reducers"], "Torch metric_reducers" - ).items() - }, - backend=_string( - policy_value.get("backend", "auto"), "Torch backend" - ), - checkpoint_owner_rank=_non_negative_integer( - policy_value.get("checkpoint_owner_rank", 0), - "Torch checkpoint_owner_rank", - ), - resume_supported=_boolean( - policy_value.get("resume_supported", True), - "Torch resume_supported", - ), - same_world_size_resume=policy_value.get("same_world_size_resume"), - rank_seeded=_boolean( - policy_value.get("rank_seeded", True), "Torch rank_seeded" - ), - checkpoint_adapter_ref=policy_value.get("checkpoint_adapter_ref"), - evidence_adapter_ref=policy_value.get("evidence_adapter_ref"), - global_loss_reducer_ref=policy_value.get("global_loss_reducer_ref"), - global_loss_reducer_api_version=policy_value.get( - "global_loss_reducer_api_version" - ), - global_loss_reducer_code_digest=policy_value.get( - "global_loss_reducer_code_digest" - ), - composite_loss_schema_id=policy_value.get( - "composite_loss_schema_id" - ), - capabilities=tuple( - _string(capability, "Torch capability") - for capability in _sequence( - policy_value.get("capabilities", ()), "Torch capabilities" - ) - ), - max_replicated_bytes_per_worker=( - _positive_integer( - policy_value["max_replicated_bytes_per_worker"], - "Torch max_replicated_bytes_per_worker", - ) - if policy_value.get("max_replicated_bytes_per_worker") - is not None - else None - ), + policy = TorchPolicy.from_dict( + { + name: value + for name, value in policy_value.items() + if name != "kind" + } ) elif kind == "framework_native": policy = FrameworkNativePolicy( diff --git a/src/tributo/algorithms/api/execution.py b/src/tributo/algorithms/api/execution.py index d5dca01..a24c36c 100644 --- a/src/tributo/algorithms/api/execution.py +++ b/src/tributo/algorithms/api/execution.py @@ -25,7 +25,6 @@ AlgorithmRequest, ) from tributo.algorithms.api.torch_runtime import ( - TorchRecoveryEnvelope, TorchStageRunIdentity, ) from tributo.util.annotations import PublicAPI @@ -85,7 +84,6 @@ class ExecutionRequest: worker_count: int resources_per_worker: WorkerResources | None = None resume_from: str | None = None - torch_recovery: TorchRecoveryEnvelope | None = None def __post_init__(self) -> None: if not isinstance(self.algorithm_request, AlgorithmRequest): @@ -113,16 +111,6 @@ def __post_init__(self) -> None: ) if self.resume_from is not None: _non_empty(self.resume_from, "resume_from") - if self.torch_recovery is not None and not isinstance( - self.torch_recovery, TorchRecoveryEnvelope - ): - raise AlgorithmConfigurationError( - "torch_recovery must be a TorchRecoveryEnvelope" - ) - if self.torch_recovery is not None and self.resume_from is not None: - raise AlgorithmConfigurationError( - "resume_from and torch_recovery are mutually exclusive" - ) @PublicAPI(stability="alpha") @@ -417,6 +405,10 @@ def __post_init__(self) -> None: raise AlgorithmConfigurationError( "Torch replicated role evidence is not identical" ) + if self.replicated_bytes_per_worker is None: + raise AlgorithmConfigurationError( + "Torch replicated role evidence requires actual bytes" + ) elif sum(rows) != self.observed_rows: raise AlgorithmConfigurationError("Torch role evidence rows do not sum") if not self.present and (self.observed_rows != 0 or any(rows)): @@ -550,6 +542,10 @@ def __post_init__(self) -> None: raise AlgorithmConfigurationError( f"Torch Stage evidence {name} is invalid" ) + if {worker.model_state_digest for worker in workers} != {self.state_digest}: + raise AlgorithmConfigurationError( + "Torch Stage worker model digests do not match its state digest" + ) if ( self.checkpoint_descriptor_digest is not None and _DIGEST.fullmatch(self.checkpoint_descriptor_digest) is None @@ -1190,7 +1186,7 @@ def kubernetes_distributed_supported(self) -> bool: def to_dict(self) -> dict[str, Any]: """Return portable receipt metadata.""" - return { + payload = { "api_version": self.api_version, "run_id": self.run_id, "plan_id": self.plan_id, @@ -1218,14 +1214,14 @@ def to_dict(self) -> dict[str, Any]: "cluster_resources": dict(sorted(self.cluster_resources.items())), "runtime_owned": self.runtime_owned, "resource_preflight": self.resource_preflight, - "torch_evidence": self.torch_evidence.to_dict() - if self.torch_evidence is not None - else None, "distributed": self.distributed, "cross_node": self.cross_node, "cluster_distributed": self.cluster_distributed, "execution_capability": self.execution_capability, } + if self.torch_evidence is not None: + payload["torch_evidence"] = self.torch_evidence.to_dict() + return payload __all__ = [ diff --git a/src/tributo/algorithms/api/models.py b/src/tributo/algorithms/api/models.py index 36f32c7..b7ea212 100644 --- a/src/tributo/algorithms/api/models.py +++ b/src/tributo/algorithms/api/models.py @@ -689,7 +689,6 @@ class RuntimeBinding: strategy: DistributionStrategy | None = None distribution_digest: str | None = None resume_from: str | None = None - torch_recovery: Mapping[str, Any] | None = None memory_bytes: int | None = None def __post_init__(self) -> None: @@ -850,24 +849,6 @@ def __post_init__(self) -> None: raise AlgorithmConfigurationError( "legacy RuntimeBinding must not carry formal resume state" ) - if self.torch_recovery is not None: - if self.strategy is not DistributionStrategy.RAY_TRAIN_TORCH: - raise AlgorithmConfigurationError( - "torch_recovery requires the Ray Train Torch runtime" - ) - from tributo.algorithms.api.torch_runtime import TorchRecoveryEnvelope - - try: - recovery = TorchRecoveryEnvelope.from_dict(self.torch_recovery) - except (TypeError, ValueError) as exc: - raise AlgorithmConfigurationError( - "runtime torch_recovery is malformed" - ) from exc - object.__setattr__(self, "torch_recovery", deep_freeze(recovery.to_dict())) - if self.resume_from is not None and self.torch_recovery is not None: - raise AlgorithmConfigurationError( - "runtime resume_from and torch_recovery are mutually exclusive" - ) @PublicAPI(stability="alpha") @@ -1725,9 +1706,6 @@ def to_dict(self, *, include_plan_id: bool = True) -> dict[str, Any]: ), "distribution_digest": self.runtime.distribution_digest, "resume_from": self.runtime.resume_from, - "torch_recovery": deep_thaw(self.runtime.torch_recovery) - if self.runtime.torch_recovery is not None - else None, } if self.runtime.memory_bytes is not None: runtime_payload["memory_bytes"] = self.runtime.memory_bytes diff --git a/src/tributo/algorithms/api/support.py b/src/tributo/algorithms/api/support.py index 614cef7..99fec10 100644 --- a/src/tributo/algorithms/api/support.py +++ b/src/tributo/algorithms/api/support.py @@ -227,9 +227,10 @@ def evidence_id(self) -> str: "issued_at": self.issued_at.isoformat(), "gate": self.gate, "result_reference": self.result_reference, - "torch_runtime_api_version": self.torch_runtime_api_version, - "torch_policy_digest": self.torch_policy_digest, } + if self.torch_runtime_api_version is not None: + payload["torch_runtime_api_version"] = self.torch_runtime_api_version + payload["torch_policy_digest"] = self.torch_policy_digest return hashlib.sha256( json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() ).hexdigest() diff --git a/src/tributo/algorithms/api/torch_runtime.py b/src/tributo/algorithms/api/torch_runtime.py index 09b2a5f..68cd411 100644 --- a/src/tributo/algorithms/api/torch_runtime.py +++ b/src/tributo/algorithms/api/torch_runtime.py @@ -12,14 +12,12 @@ import math import os import re -import threading from collections.abc import Mapping from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Generator, Protocol, cast, runtime_checkable -from urllib.parse import urlsplit +from typing import TYPE_CHECKING, Any, Generator, Protocol, runtime_checkable from tributo._common.immutable import deep_freeze from tributo.algorithms.api.errors import ( @@ -29,7 +27,6 @@ from tributo.util.annotations import PublicAPI if TYPE_CHECKING: - from tributo.algorithms.spi.execution import RuntimeExecutionEnvelope from tributo.algorithms.spi.torch import TorchStageContext @@ -233,11 +230,10 @@ def identity_digest(self) -> str: @property def run_config_name(self) -> str: - return torch_run_config_name(self) + return _torch_run_config_name(self) -@PublicAPI(stability="alpha") -def torch_run_config_name(identity: TorchStageRunIdentity) -> str: +def _torch_run_config_name(identity: TorchStageRunIdentity) -> str: """Build the Core-owned deterministic Ray ``RunConfig.name``.""" if not isinstance(identity, TorchStageRunIdentity): raise AlgorithmConfigurationError("Torch run identity is required") @@ -255,141 +251,6 @@ def torch_run_config_name(identity: TorchStageRunIdentity) -> str: return name -@PublicAPI(stability="alpha") -def claim_torch_run_directory( - storage_path: str | os.PathLike[str], - identity: TorchStageRunIdentity, - *, - status: str = "running", -) -> Path: - """Atomically claim a Stage directory and its identity manifest. - - Existing directories are reusable only for the same identity and a - retryable status. Nothing is overwritten or silently renamed. - """ - if status not in {"created", "running", "failed_retryable"}: - raise AlgorithmConfigurationError("Torch run directory status is not retryable") - root = Path(storage_path) - raw_storage = str(storage_path) - if "://" in raw_storage and not raw_storage.startswith("file://"): - # Use the filesystem abstraction only for the identity manifest. The - # actual Ray checkpoint payload remains owned by Ray's storage backend. - try: - import pyarrow.fs as pafs - - filesystem, prefix = pafs.FileSystem.from_uri(raw_storage) - run_path = f"{prefix.rstrip('/')}/{torch_run_config_name(identity)}" - manifest_path = f"{run_path}/torch_identity_manifest.json" - filesystem.create_dir(run_path, recursive=True) - info = filesystem.get_file_info(manifest_path) - payload = { - "schema_version": 1, - "snapshot_schema": 1, - "identity": identity.to_dict(), - "status": status, - "stale": False, - } - if info.type is pafs.FileType.File: - existing = json.loads( - filesystem.open_input_file(manifest_path).read().decode("utf-8") - ) - if ( - not isinstance(existing, Mapping) - or existing.get("identity") != identity.to_dict() - or existing.get("snapshot_schema", 1) != 1 - or existing.get("stale", False) - or existing.get("status") - not in {"created", "running", "failed_retryable"} - ): - raise AlgorithmExecutionError( - "Torch remote run directory identity collision or stale snapshot" - ) - elif info.type is pafs.FileType.NotFound: - with filesystem.open_output_stream(manifest_path) as stream: - stream.write((_canonical_json(payload) + "\n").encode("utf-8")) - else: - raise AlgorithmExecutionError( - "Torch remote run identity manifest has an invalid type" - ) - return Path(f"{raw_storage.rstrip('/')}/{torch_run_config_name(identity)}") - except AlgorithmExecutionError: - raise - except (ImportError, OSError, TypeError, ValueError) as exc: - raise AlgorithmExecutionError( - "failed to claim remote Torch run identity manifest" - ) from exc - if not root.is_absolute(): - raise AlgorithmConfigurationError("Torch storage_path must be absolute") - directory = root / torch_run_config_name(identity) - manifest = directory / "torch_identity_manifest.json" - try: - directory.mkdir(parents=True, exist_ok=False) - except FileExistsError: - if not manifest.is_file(): - raise AlgorithmExecutionError( - "Torch run directory has no identity manifest" - ) from None - try: - existing = json.loads(manifest.read_text(encoding="utf-8")) - except (OSError, TypeError, ValueError) as exc: - raise AlgorithmExecutionError( - "Torch run identity manifest is damaged" - ) from exc - if ( - not isinstance(existing, Mapping) - or existing.get("identity") != identity.to_dict() - ): - raise AlgorithmExecutionError( - "Torch run directory identity collision" - ) from None - if existing.get("snapshot_schema", 1) != 1 or existing.get("stale", False): - raise AlgorithmExecutionError( - "Torch run directory contains a stale snapshot" - ) from None - if existing.get("status") not in {"created", "running", "failed_retryable"}: - raise AlgorithmExecutionError( - "Torch run directory is terminal or stale" - ) from None - return directory - payload = { - "schema_version": 1, - "snapshot_schema": 1, - "identity": identity.to_dict(), - "status": status, - "stale": False, - } - try: - fd = os.open(manifest, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) - with os.fdopen(fd, "w", encoding="utf-8") as stream: - stream.write(_canonical_json(payload)) - stream.write("\n") - except OSError as exc: - raise AlgorithmExecutionError( - "failed to atomically write Torch identity manifest" - ) from exc - return directory - - -@PublicAPI(stability="alpha") -def validate_torch_retry_identity( - descriptor: "TorchCheckpointDescriptor", - identity: TorchStageRunIdentity, - *, - world_size: int, -) -> None: - """Reject stale Ray retry snapshots before model state is deserialized.""" - if not isinstance(descriptor, TorchCheckpointDescriptor): - raise AlgorithmExecutionError("retry checkpoint has no Torch descriptor") - if descriptor.identity.to_dict() != identity.to_dict(): - raise AlgorithmExecutionError( - "retry checkpoint identity does not match current Stage" - ) - if descriptor.world_size != world_size: - raise AlgorithmExecutionError( - "retry checkpoint world size does not match current Stage" - ) - - @PublicAPI(stability="alpha") @dataclass(frozen=True) class TorchLossContribution: @@ -759,72 +620,23 @@ def reduce( @PublicAPI(stability="alpha") @dataclass(frozen=True) -class TorchPreflightTokenData: - """Immutable identity data produced by Torch preflight.""" +class TorchCheckpointPayloadDraft: + """Local payload directory reported through the Core Ray Train helper.""" - run_id: str - invocation_id: str - algorithm: str - implementation_ref: str - implementation_code_digest: str - policy_digest: str - execution_plan_digest: str - runtime_id: str - reducer_identity: str | None = None - plan_digest: str | None = None + checkpoint_dir: str | os.PathLike[str] + checkpoint_owner_rank: int = 0 def __post_init__(self) -> None: - for name in ( - "run_id", - "invocation_id", - "algorithm", - "implementation_ref", - "runtime_id", - ): - if not isinstance(getattr(self, name), str) or not getattr(self, name): - raise AlgorithmConfigurationError(f"{name} must be non-empty") - _digest_value(self.implementation_code_digest, "implementation_code_digest") - _digest_value(self.policy_digest, "policy_digest") - _digest_value(self.execution_plan_digest, "execution_plan_digest") - if self.reducer_identity is not None and ( - not isinstance(self.reducer_identity, str) - or ":" not in self.reducer_identity + if not isinstance(self.checkpoint_dir, (str, os.PathLike)): + raise AlgorithmConfigurationError("checkpoint_dir must be path-like") + if ( + not isinstance(self.checkpoint_owner_rank, int) + or isinstance(self.checkpoint_owner_rank, bool) + or self.checkpoint_owner_rank < 0 ): raise AlgorithmConfigurationError( - "reducer_identity must be qualified when provided" + "checkpoint_owner_rank must be non-negative" ) - if self.plan_digest is not None: - _digest_value(self.plan_digest, "plan_digest") - - def to_dict(self) -> dict[str, Any]: - return { - "run_id": self.run_id, - "invocation_id": self.invocation_id, - "algorithm": self.algorithm, - "implementation_ref": self.implementation_ref, - "implementation_code_digest": self.implementation_code_digest, - "policy_digest": self.policy_digest, - "execution_plan_digest": self.execution_plan_digest, - "runtime_id": self.runtime_id, - "reducer_identity": self.reducer_identity, - "plan_digest": self.plan_digest, - } - - -@PublicAPI(stability="alpha") -@runtime_checkable -class TorchCheckpointPayloadDraft(Protocol): - """Optional Core payload draft hook used by checkpoint report tests.""" - - checkpoint_dir: str | os.PathLike[str] - - def report( - self, - *, - metrics: Mapping[str, object], - stage_context: object, - completed_step: int, - ) -> None: ... @PublicAPI(stability="alpha") @@ -879,222 +691,6 @@ def __exit__(self, exc_type: object, exc: object, traceback: object) -> None: self.close() -@PublicAPI(stability="alpha") -class TorchPreflightLease: - """Invocation-local one-shot ownership token for a preflight result.""" - - __slots__ = ("_data", "_state", "_lock") - - def __init__(self, data: TorchPreflightTokenData) -> None: - if not isinstance(data, TorchPreflightTokenData): - raise AlgorithmConfigurationError("preflight lease requires TokenData") - self._data = data - self._state = "fresh" - self._lock = threading.Lock() - - @property - def state(self) -> str: - with self._lock: - return self._state - - def _matches( - self, *, run_id: str, invocation_id: str, plan_digest: str, runtime_id: str - ) -> bool: - return ( - self._data.run_id == run_id - and self._data.invocation_id == invocation_id - and ( - self._data.plan_digest is None - and self._data.execution_plan_digest == plan_digest - or self._data.plan_digest is not None - and self._data.plan_digest == plan_digest - ) - and self._data.runtime_id == runtime_id - ) - - def claim( - self, - *, - run_id: str, - invocation_id: str, - plan_digest: str, - runtime_id: str, - ) -> None: - with self._lock: - if self._state != "fresh": - raise AlgorithmExecutionError("Torch preflight lease is not fresh") - if not self._matches( - run_id=run_id, - invocation_id=invocation_id, - plan_digest=plan_digest, - runtime_id=runtime_id, - ): - raise AlgorithmExecutionError("Torch preflight lease identity mismatch") - self._state = "claimed" - - def consume( - self, - *, - run_id: str, - invocation_id: str, - plan_digest: str, - runtime_id: str, - ) -> TorchPreflightTokenData: - with self._lock: - if self._state != "claimed": - raise AlgorithmExecutionError("Torch preflight lease is not claimed") - if not self._matches( - run_id=run_id, - invocation_id=invocation_id, - plan_digest=plan_digest, - runtime_id=runtime_id, - ): - raise AlgorithmExecutionError("Torch preflight token identity mismatch") - self._state = "consumed" - return self._data - - def close(self) -> None: - with self._lock: - if self._state in {"fresh", "claimed"}: - self._state = "closed" - - @property - def data(self) -> TorchPreflightTokenData: - with self._lock: - if self._state not in {"fresh", "claimed"}: - raise AlgorithmExecutionError("Torch preflight lease is closed") - return self._data - - def __copy__(self) -> object: - raise TypeError("TorchPreflightLease cannot be copied") - - def __deepcopy__(self, memo: dict[int, object]) -> object: - del memo - raise TypeError("TorchPreflightLease cannot be copied") - - def __getstate__(self) -> object: - raise TypeError("TorchPreflightLease cannot be serialized") - - -@PublicAPI(stability="alpha") -@dataclass(frozen=True) -class TorchWorkerControlEnvelope: - """Credential-free serialized control for initial Stage Checkpoint input.""" - - schema_version: int - run_id: str - invocation_id: str - source_stage_id: str | None - target_stage_id: str - purpose: str - checkpoint_locator: "TorchCheckpointLocator" - checkpoint_descriptor_digest: str - policy_digest: str - execution_plan_digest: str - - def __post_init__(self) -> None: - if self.schema_version != 1: - raise AlgorithmConfigurationError( - "unsupported Torch worker control version" - ) - if self.purpose not in {"stage_dependency", "cross_run_initial_recovery"}: - raise AlgorithmConfigurationError("invalid Torch worker control purpose") - if self.purpose == "stage_dependency" and ( - not isinstance(self.source_stage_id, str) or not self.source_stage_id - ): - raise AlgorithmConfigurationError( - "stage_dependency control requires a source_stage_id" - ) - for name in ("run_id", "invocation_id", "target_stage_id"): - if not isinstance(getattr(self, name), str) or not getattr(self, name): - raise AlgorithmConfigurationError(f"{name} must be non-empty") - if self.source_stage_id == self.target_stage_id: - raise AlgorithmConfigurationError( - "Torch control source and target Stage must differ" - ) - if not isinstance(self.checkpoint_locator, TorchCheckpointLocator): - raise AlgorithmConfigurationError("Torch checkpoint locator is invalid") - _digest_value(self.checkpoint_descriptor_digest, "checkpoint_descriptor_digest") - if ( - self.checkpoint_locator.descriptor_digest - != self.checkpoint_descriptor_digest - ): - raise AlgorithmConfigurationError( - "Torch control locator and descriptor digests do not match" - ) - _digest_value(self.policy_digest, "policy_digest") - _digest_value(self.execution_plan_digest, "execution_plan_digest") - - def to_dict(self) -> dict[str, Any]: - return { - "schema_version": self.schema_version, - "run_id": self.run_id, - "invocation_id": self.invocation_id, - "source_stage_id": self.source_stage_id, - "target_stage_id": self.target_stage_id, - "purpose": self.purpose, - "checkpoint_locator": self.checkpoint_locator.to_dict(), - "checkpoint_descriptor_digest": self.checkpoint_descriptor_digest, - "policy_digest": self.policy_digest, - "execution_plan_digest": self.execution_plan_digest, - } - - @classmethod - def from_dict(cls, value: Mapping[str, Any]) -> "TorchWorkerControlEnvelope": - raw_locator = value.get("checkpoint_locator") - if not isinstance(raw_locator, Mapping): - raise AlgorithmConfigurationError( - "Torch worker control locator must be a typed locator" - ) - locator = TorchCheckpointLocator.from_dict(raw_locator) - try: - return cls( - schema_version=value["schema_version"], - run_id=value["run_id"], - invocation_id=value["invocation_id"], - source_stage_id=value.get("source_stage_id"), - target_stage_id=value["target_stage_id"], - purpose=value["purpose"], - checkpoint_locator=locator, - checkpoint_descriptor_digest=value["checkpoint_descriptor_digest"], - policy_digest=value["policy_digest"], - execution_plan_digest=value["execution_plan_digest"], - ) - except KeyError as exc: - raise AlgorithmConfigurationError( - "Torch worker control is missing a field" - ) from exc - - -@PublicAPI(stability="alpha") -@dataclass(frozen=True) -class TorchRuntimeExecutionEnvelope: - """Driver-local Torch envelope carrying a claimed preflight lease.""" - - base: "RuntimeExecutionEnvelope" - preflight_lease: TorchPreflightLease - - def __post_init__(self) -> None: - if not hasattr(self.base, "plan") or not hasattr(self.base, "run_id"): - raise AlgorithmConfigurationError( - "Torch envelope requires a RuntimeExecutionEnvelope" - ) - if not isinstance(self.preflight_lease, TorchPreflightLease): - raise AlgorithmConfigurationError( - "Torch envelope requires a preflight lease" - ) - - def __getstate__(self) -> object: - raise TypeError("Torch runtime envelope cannot be serialized") - - def __copy__(self) -> object: - raise TypeError("Torch runtime envelope cannot be copied") - - def __deepcopy__(self, memo: dict[int, object]) -> object: - del memo - raise TypeError("Torch runtime envelope cannot be copied") - - def _loss_numerator_and_normalizer(loss: TorchStepLoss) -> tuple[object, float]: if isinstance(loss, TorchLossContribution): return loss.numerator, loss.normalizer @@ -1259,13 +855,9 @@ def validate_metric_metadata(value: object, path: str = "metrics") -> None: ) validate_metric_metadata(metrics) - reporter = getattr(payload_draft, "report", None) - checkpoint_dir = getattr(payload_draft, "checkpoint_dir", None) - if not callable(reporter) or checkpoint_dir is None: - raise AlgorithmExecutionError( - "Torch checkpoints require a Core payload draft with checkpoint_dir and report()" - ) - root = Path(checkpoint_dir) + if not isinstance(payload_draft, TorchCheckpointPayloadDraft): + raise AlgorithmExecutionError("Torch checkpoints require a typed payload draft") + root = Path(payload_draft.checkpoint_dir) if root.is_symlink() or not root.is_dir(): raise AlgorithmExecutionError("Torch checkpoint payload directory is missing") root = root.resolve() @@ -1277,6 +869,10 @@ def validate_metric_metadata(value: object, path: str = "metrics") -> None: ) if runtime is None: raise AlgorithmExecutionError("Torch checkpoint Stage context has no runtime") + if payload_draft.checkpoint_owner_rank >= runtime.world_size: + raise AlgorithmExecutionError( + "Torch checkpoint owner rank is outside the Worker group" + ) binding_digest = getattr(runtime, "input_binding_digest", None) if not isinstance(binding_digest, str) or len(binding_digest) != 64: raise AlgorithmExecutionError( @@ -1363,8 +959,7 @@ def validate_metric_metadata(value: object, path: str = "metrics") -> None: implementation_code_digest=identity.implementation_code_digest, payload_files=dict(files), adapter_identity=getattr(runtime, "adapter_identity", None), - resume_supported=getattr(runtime, "resume_supported", True), - same_world_size_resume=getattr(runtime, "same_world_size_resume", True), + resume_supported=getattr(runtime, "resume_supported", False), ) temporary_descriptor = root / ".torch_checkpoint_descriptor.tmp" if temporary_descriptor.exists() or temporary_descriptor.is_symlink(): @@ -1391,11 +986,16 @@ def validate_metric_metadata(value: object, path: str = "metrics") -> None: ) from exc report_metrics = dict(metrics) report_metrics["checkpoint_descriptor"] = descriptor.to_dict() - reporter( - metrics=report_metrics, - stage_context=stage_context, - completed_step=completed_step, + import ray.train + from ray.train import Checkpoint + + rank = ray.train.get_context().get_world_rank() + checkpoint = ( + Checkpoint.from_directory(str(root)) + if rank == payload_draft.checkpoint_owner_rank + else None ) + ray.train.report(report_metrics, checkpoint=checkpoint) def _scan_checkpoint_files(root: Path) -> dict[str, str]: @@ -1406,7 +1006,6 @@ def _scan_checkpoint_files(root: Path) -> dict[str, str]: raise AlgorithmExecutionError("Torch checkpoint payload escapes its root") if path.name in { "torch_checkpoint_descriptor.json", - "torch_stage_commit.json", ".metadata.json", }: continue @@ -1416,8 +1015,7 @@ def _scan_checkpoint_files(root: Path) -> dict[str, str]: return files -@PublicAPI(stability="alpha") -def describe_torch_checkpoint( +def _describe_torch_checkpoint( checkpoint_ref: TorchCheckpointRef, checkpoint_context: object, ) -> TorchCheckpointDescriptor: @@ -1442,25 +1040,6 @@ def describe_torch_checkpoint( descriptor = parsed if descriptor.run_config_name != descriptor.identity.run_config_name: raise AlgorithmExecutionError("Torch checkpoint RunConfig name drifted") - commit_path = root / "torch_stage_commit.json" - if commit_path.exists() or commit_path.is_symlink(): - if commit_path.is_symlink() or not commit_path.is_file(): - raise AlgorithmExecutionError( - "Torch checkpoint commit marker is invalid" - ) - try: - commit = json.loads(commit_path.read_text(encoding="utf-8")) - except (OSError, TypeError, ValueError) as exc: - raise AlgorithmExecutionError( - "Torch checkpoint commit marker is malformed" - ) from exc - if not isinstance(commit, Mapping) or ( - commit.get("identity") != descriptor.identity.to_dict() - or commit.get("descriptor_digest") != descriptor.digest - ): - raise AlgorithmExecutionError( - "Torch checkpoint commit marker does not match descriptor" - ) files = _scan_checkpoint_files(root) if dict(descriptor.payload_files) != files: raise AlgorithmExecutionError( @@ -1497,7 +1076,7 @@ def describe_torch_checkpoint( raise AlgorithmExecutionError( "Torch checkpoint descriptor Adapter identity drifted" ) - for name in ("state_layout", "resume_supported", "same_world_size_resume"): + for name in ("state_layout", "resume_supported"): expected = getattr(runtime, name, None) if getattr(descriptor, name) != expected: raise AlgorithmExecutionError( @@ -1529,529 +1108,6 @@ def _opened_checkpoint(checkpoint: object) -> Generator[Path, None, None]: raise AlgorithmExecutionError("Torch checkpoint could not be opened") from exc -@PublicAPI(stability="alpha") -@dataclass(frozen=True) -class TorchCheckpointLocator: - """Credential-free persistent location for a validated Torch checkpoint.""" - - uri: str - descriptor_digest: str - schema_version: int = 1 - - def __post_init__(self) -> None: - if self.schema_version != 1: - raise AlgorithmConfigurationError("unsupported Torch locator version") - if not isinstance(self.uri, str) or not self.uri or "\x00" in self.uri: - raise AlgorithmConfigurationError("Torch checkpoint locator URI is invalid") - if self.uri.startswith(("/", "file://")): - raise AlgorithmConfigurationError( - "Torch checkpoint locator must not persist a local path" - ) - try: - parsed = urlsplit(self.uri) - except ValueError as exc: - raise AlgorithmConfigurationError( - "Torch checkpoint locator URI is invalid" - ) from exc - if parsed.username is not None or parsed.password is not None: - raise AlgorithmConfigurationError( - "Torch checkpoint locator must not contain URI userinfo" - ) - if parsed.query or parsed.fragment: - raise AlgorithmConfigurationError( - "Torch checkpoint locator must not contain query or fragment data" - ) - _digest_value(self.descriptor_digest, "descriptor_digest") - - def to_dict(self) -> dict[str, Any]: - return { - "schema_version": self.schema_version, - "uri": self.uri, - "descriptor_digest": self.descriptor_digest, - } - - @classmethod - def from_dict(cls, value: Mapping[str, Any]) -> "TorchCheckpointLocator": - try: - return cls( - uri=value["uri"], - descriptor_digest=value["descriptor_digest"], - schema_version=value.get("schema_version", 1), - ) - except KeyError as exc: - raise AlgorithmConfigurationError( - "Torch checkpoint locator is incomplete" - ) from exc - - -@PublicAPI(stability="alpha") -@dataclass(frozen=True) -class TorchRecoveryEnvelope: - """Credential-free cross-Run recovery state for one Torch execution plan.""" - - completed_stage_ids: tuple[str, ...] = () - stage_checkpoints: Mapping[str, TorchCheckpointLocator] = field( - default_factory=dict - ) - active_stage_id: str | None = None - active_checkpoint: TorchCheckpointLocator | None = None - schema_version: int = 1 - - def __post_init__(self) -> None: - if self.schema_version != 1: - raise AlgorithmConfigurationError( - "unsupported Torch recovery envelope version" - ) - completed = tuple(self.completed_stage_ids) - if any(not isinstance(stage, str) or not stage for stage in completed): - raise AlgorithmConfigurationError( - "Torch recovery completed_stage_ids must be non-empty strings" - ) - if len(set(completed)) != len(completed): - raise AlgorithmConfigurationError( - "Torch recovery completed_stage_ids must be unique" - ) - checkpoints: dict[str, TorchCheckpointLocator] = {} - for stage_id, locator in self.stage_checkpoints.items(): - if not isinstance(stage_id, str) or not stage_id: - raise AlgorithmConfigurationError( - "Torch recovery checkpoint stage IDs must be non-empty" - ) - if not isinstance(locator, TorchCheckpointLocator): - raise AlgorithmConfigurationError( - "Torch recovery stage checkpoints must use locators" - ) - checkpoints[stage_id] = locator - if set(checkpoints) != set(completed): - raise AlgorithmConfigurationError( - "Torch recovery stage checkpoints must match completed stages" - ) - if self.active_stage_id is None: - if self.active_checkpoint is not None: - raise AlgorithmConfigurationError( - "Torch recovery active checkpoint requires active_stage_id" - ) - else: - if not isinstance(self.active_stage_id, str) or not self.active_stage_id: - raise AlgorithmConfigurationError( - "Torch recovery active_stage_id must be non-empty" - ) - if self.active_stage_id in completed: - raise AlgorithmConfigurationError( - "Torch recovery active Stage cannot be completed" - ) - if not isinstance(self.active_checkpoint, TorchCheckpointLocator): - raise AlgorithmConfigurationError( - "Torch recovery active Stage requires a checkpoint locator" - ) - object.__setattr__(self, "completed_stage_ids", completed) - object.__setattr__( - self, - "stage_checkpoints", - MappingProxyType(dict(sorted(checkpoints.items()))), - ) - - def to_dict(self) -> dict[str, Any]: - return { - "schema_version": self.schema_version, - "completed_stage_ids": list(self.completed_stage_ids), - "stage_checkpoints": { - stage_id: locator.to_dict() - for stage_id, locator in self.stage_checkpoints.items() - }, - "active_stage_id": self.active_stage_id, - "active_checkpoint": ( - self.active_checkpoint.to_dict() - if self.active_checkpoint is not None - else None - ), - } - - @classmethod - def from_dict(cls, value: Mapping[str, Any]) -> "TorchRecoveryEnvelope": - raw_checkpoints = value.get("stage_checkpoints", {}) - if not isinstance(raw_checkpoints, Mapping): - raise AlgorithmConfigurationError( - "Torch recovery stage_checkpoints must be a mapping" - ) - raw_active = value.get("active_checkpoint") - parsed_checkpoints: dict[str, TorchCheckpointLocator] = {} - for stage_id, locator in raw_checkpoints.items(): - if not isinstance(stage_id, str) or not isinstance(locator, Mapping): - raise AlgorithmConfigurationError( - "Torch recovery stage_checkpoints entries are malformed" - ) - parsed_checkpoints[stage_id] = TorchCheckpointLocator.from_dict(locator) - try: - return cls( - schema_version=value.get("schema_version", 1), - completed_stage_ids=tuple(value.get("completed_stage_ids", ())), - stage_checkpoints=parsed_checkpoints, - active_stage_id=cast(str | None, value.get("active_stage_id")), - active_checkpoint=( - TorchCheckpointLocator.from_dict(raw_active) - if isinstance(raw_active, Mapping) - else None - ), - ) - except (KeyError, TypeError, ValueError) as exc: - raise AlgorithmConfigurationError( - "Torch recovery envelope is malformed" - ) from exc - - -@PublicAPI(stability="alpha") -@dataclass(frozen=True) -class TorchRankProgressStatistics: - """Typed per-rank prefix statistics restored with a Torch checkpoint.""" - - rows_processed: int = 0 - coverage_totals: Mapping[str, int] = field(default_factory=dict) - loss_numerator_total: float = 0.0 - loss_normalizer_total: float = 0.0 - metric_totals: Mapping[str, tuple[float, float]] = field(default_factory=dict) - evaluation_totals: Mapping[str, tuple[float, float]] = field(default_factory=dict) - reducer_observation: Mapping[str, object] = field(default_factory=dict) - - def __post_init__(self) -> None: - if ( - not isinstance(self.rows_processed, int) - or isinstance(self.rows_processed, bool) - or self.rows_processed < 0 - ): - raise AlgorithmConfigurationError("Torch rank rows_processed is invalid") - coverage = dict(self.coverage_totals) - if any( - not isinstance(name, str) - or not name - or not isinstance(value, int) - or isinstance(value, bool) - or value < 0 - for name, value in coverage.items() - ): - raise AlgorithmConfigurationError("Torch rank coverage totals are invalid") - - def normalize_totals( - values: Mapping[str, tuple[float, float]], - ) -> dict[str, tuple[float, float]]: - normalized: dict[str, tuple[float, float]] = {} - for name, pair in values.items(): - if ( - not isinstance(name, str) - or not name - or not isinstance(pair, (list, tuple)) - or len(pair) != 2 - ): - raise AlgorithmConfigurationError( - "Torch rank metric totals are invalid" - ) - numerator = _finite_number(pair[0], f"rank metric[{name}] numerator") - normalizer = _finite_number(pair[1], f"rank metric[{name}] normalizer") - if normalizer < 0: - raise AlgorithmConfigurationError( - "Torch rank metric normalizer is invalid" - ) - normalized[name] = (numerator, normalizer) - return normalized - - loss_numerator = _finite_number( - self.loss_numerator_total, "rank loss numerator total" - ) - loss_normalizer = _finite_number( - self.loss_normalizer_total, "rank loss normalizer total" - ) - if loss_normalizer < 0: - raise AlgorithmConfigurationError("Torch rank loss normalizer is invalid") - metric_totals = normalize_totals(self.metric_totals) - evaluation_totals = normalize_totals(self.evaluation_totals) - _validate_bounded_evidence( - self.reducer_observation, path="rank.reducer_observation" - ) - object.__setattr__( - self, "coverage_totals", MappingProxyType(dict(sorted(coverage.items()))) - ) - object.__setattr__(self, "loss_numerator_total", loss_numerator) - object.__setattr__(self, "loss_normalizer_total", loss_normalizer) - object.__setattr__( - self, "metric_totals", MappingProxyType(dict(sorted(metric_totals.items()))) - ) - object.__setattr__( - self, - "evaluation_totals", - MappingProxyType(dict(sorted(evaluation_totals.items()))), - ) - object.__setattr__( - self, "reducer_observation", deep_freeze(self.reducer_observation) - ) - - def to_dict(self) -> dict[str, Any]: - return { - "rows_processed": self.rows_processed, - "coverage_totals": dict(self.coverage_totals), - "loss_numerator_total": self.loss_numerator_total, - "loss_normalizer_total": self.loss_normalizer_total, - "metric_totals": { - name: list(pair) for name, pair in self.metric_totals.items() - }, - "evaluation_totals": { - name: list(pair) for name, pair in self.evaluation_totals.items() - }, - "reducer_observation": dict(self.reducer_observation), - } - - @classmethod - def from_dict(cls, value: Mapping[str, Any]) -> "TorchRankProgressStatistics": - try: - return cls( - rows_processed=value.get("rows_processed", 0), - coverage_totals=value.get("coverage_totals", {}), - loss_numerator_total=value.get("loss_numerator_total", 0.0), - loss_normalizer_total=value.get("loss_normalizer_total", 0.0), - metric_totals=value.get("metric_totals", {}), - evaluation_totals=value.get("evaluation_totals", {}), - reducer_observation=value.get("reducer_observation", {}), - ) - except (TypeError, ValueError) as exc: - raise AlgorithmConfigurationError( - "Torch rank progress statistics are malformed" - ) from exc - - -@PublicAPI(stability="alpha") -@dataclass(frozen=True) -class TorchCheckpointProgress: - """Deterministic cursor state required for exact Torch recovery.""" - - epoch: int - micro_batch_cursor: int - optimizer_step: int - scheduler_step: int - accumulation_steps: int - dataset_cursor_by_rank: Mapping[str, int] = field(default_factory=dict) - shuffle_seed: int | None = None - rows_processed: int = 0 - coverage_totals: Mapping[str, int] = field(default_factory=dict) - loss_numerator_total: float = 0.0 - loss_normalizer_total: float = 0.0 - metric_totals: Mapping[str, tuple[float, float]] = field(default_factory=dict) - evaluation_totals: Mapping[str, tuple[float, float]] = field(default_factory=dict) - rank_statistics: Mapping[str, TorchRankProgressStatistics] = field( - default_factory=dict - ) - epoch_scheduler_applied: bool = False - schema_version: int = 1 - - def __post_init__(self) -> None: - if self.schema_version != 1: - raise AlgorithmConfigurationError( - "unsupported Torch checkpoint progress version" - ) - for name in ( - "epoch", - "micro_batch_cursor", - "optimizer_step", - "scheduler_step", - "accumulation_steps", - ): - value = getattr(self, name) - if ( - not isinstance(value, int) - or isinstance(value, bool) - or value < 0 - or (name == "accumulation_steps" and value < 1) - ): - raise AlgorithmConfigurationError( - f"Torch checkpoint progress {name} is invalid" - ) - cursors: dict[str, int] = {} - for rank, cursor in self.dataset_cursor_by_rank.items(): - if ( - not isinstance(rank, str) - or not rank - or not isinstance(cursor, int) - or isinstance(cursor, bool) - or cursor < 0 - ): - raise AlgorithmConfigurationError( - "Torch checkpoint dataset cursors are invalid" - ) - cursors[rank] = cursor - if self.shuffle_seed is not None and ( - not isinstance(self.shuffle_seed, int) - or isinstance(self.shuffle_seed, bool) - ): - raise AlgorithmConfigurationError( - "Torch checkpoint shuffle_seed is invalid" - ) - if ( - not isinstance(self.rows_processed, int) - or isinstance(self.rows_processed, bool) - or self.rows_processed < 0 - ): - raise AlgorithmConfigurationError( - "Torch checkpoint rows_processed is invalid" - ) - coverage = dict(self.coverage_totals) - if any( - not isinstance(name, str) - or not name - or not isinstance(value, int) - or isinstance(value, bool) - or value < 0 - for name, value in coverage.items() - ): - raise AlgorithmConfigurationError( - "Torch checkpoint coverage totals are invalid" - ) - loss_numerator = _finite_number( - self.loss_numerator_total, "checkpoint loss numerator total" - ) - loss_normalizer = _finite_number( - self.loss_normalizer_total, "checkpoint loss normalizer total" - ) - if loss_normalizer < 0: - raise AlgorithmConfigurationError( - "Torch checkpoint loss normalizer is invalid" - ) - - def normalize_totals( - values: Mapping[str, tuple[float, float]], - ) -> dict[str, tuple[float, float]]: - normalized: dict[str, tuple[float, float]] = {} - for name, pair in values.items(): - if ( - not isinstance(name, str) - or not name - or not isinstance(pair, (list, tuple)) - or len(pair) != 2 - ): - raise AlgorithmConfigurationError( - "Torch checkpoint metric totals are invalid" - ) - numerator = _finite_number( - pair[0], f"checkpoint metric[{name}] numerator" - ) - normalizer = _finite_number( - pair[1], f"checkpoint metric[{name}] normalizer" - ) - if normalizer < 0: - raise AlgorithmConfigurationError( - "Torch checkpoint metric normalizer is invalid" - ) - normalized[name] = (numerator, normalizer) - return normalized - - metric_totals = normalize_totals(self.metric_totals) - evaluation_totals = normalize_totals(self.evaluation_totals) - rank_statistics = dict(self.rank_statistics) - if len(rank_statistics) > 1024 or any( - not isinstance(rank, str) - or not rank - or not isinstance(stats, TorchRankProgressStatistics) - for rank, stats in rank_statistics.items() - ): - raise AlgorithmConfigurationError( - "Torch checkpoint rank statistics are invalid" - ) - if not isinstance(self.epoch_scheduler_applied, bool): - raise AlgorithmConfigurationError( - "Torch checkpoint epoch_scheduler_applied must be boolean" - ) - object.__setattr__( - self, - "dataset_cursor_by_rank", - MappingProxyType(dict(sorted(cursors.items()))), - ) - object.__setattr__(self, "rows_processed", self.rows_processed) - object.__setattr__( - self, "coverage_totals", MappingProxyType(dict(sorted(coverage.items()))) - ) - object.__setattr__(self, "loss_numerator_total", loss_numerator) - object.__setattr__(self, "loss_normalizer_total", loss_normalizer) - object.__setattr__( - self, "metric_totals", MappingProxyType(dict(sorted(metric_totals.items()))) - ) - object.__setattr__( - self, - "evaluation_totals", - MappingProxyType(dict(sorted(evaluation_totals.items()))), - ) - object.__setattr__( - self, - "rank_statistics", - MappingProxyType(dict(sorted(rank_statistics.items()))), - ) - - def to_dict(self) -> dict[str, Any]: - return { - "schema_version": self.schema_version, - "epoch": self.epoch, - "micro_batch_cursor": self.micro_batch_cursor, - "optimizer_step": self.optimizer_step, - "scheduler_step": self.scheduler_step, - "accumulation_steps": self.accumulation_steps, - "dataset_cursor_by_rank": dict(self.dataset_cursor_by_rank), - "shuffle_seed": self.shuffle_seed, - "rows_processed": self.rows_processed, - "coverage_totals": dict(self.coverage_totals), - "loss_numerator_total": self.loss_numerator_total, - "loss_normalizer_total": self.loss_normalizer_total, - "metric_totals": { - name: list(pair) for name, pair in self.metric_totals.items() - }, - "evaluation_totals": { - name: list(pair) for name, pair in self.evaluation_totals.items() - }, - "rank_statistics": { - rank: statistics.to_dict() - for rank, statistics in self.rank_statistics.items() - }, - "epoch_scheduler_applied": self.epoch_scheduler_applied, - } - - @classmethod - def from_dict(cls, value: Mapping[str, Any]) -> "TorchCheckpointProgress": - raw_cursors = value.get("dataset_cursor_by_rank", {}) - if not isinstance(raw_cursors, Mapping): - raise AlgorithmConfigurationError( - "Torch checkpoint dataset cursors must be a mapping" - ) - raw_statistics = value.get("rank_statistics", {}) - if not isinstance(raw_statistics, Mapping) or any( - not isinstance(rank, str) or not isinstance(statistics, Mapping) - for rank, statistics in raw_statistics.items() - ): - raise AlgorithmConfigurationError( - "Torch checkpoint rank statistics must be a typed mapping" - ) - try: - return cls( - schema_version=value.get("schema_version", 1), - epoch=value["epoch"], - micro_batch_cursor=value["micro_batch_cursor"], - optimizer_step=value["optimizer_step"], - scheduler_step=value["scheduler_step"], - accumulation_steps=value["accumulation_steps"], - dataset_cursor_by_rank=dict(raw_cursors), - shuffle_seed=value.get("shuffle_seed"), - rows_processed=value.get("rows_processed", 0), - coverage_totals=value.get("coverage_totals", {}), - loss_numerator_total=value.get("loss_numerator_total", 0.0), - loss_normalizer_total=value.get("loss_normalizer_total", 0.0), - metric_totals=value.get("metric_totals", {}), - evaluation_totals=value.get("evaluation_totals", {}), - rank_statistics={ - rank: TorchRankProgressStatistics.from_dict(statistics) - for rank, statistics in raw_statistics.items() - }, - epoch_scheduler_applied=value.get("epoch_scheduler_applied", False), - ) - except (KeyError, TypeError, ValueError) as exc: - raise AlgorithmConfigurationError( - "Torch checkpoint progress is malformed" - ) from exc - - @PublicAPI(stability="alpha") @dataclass(frozen=True) class TorchCheckpointDescriptor: @@ -2069,8 +1125,7 @@ class TorchCheckpointDescriptor: implementation_code_digest: str payload_files: Mapping[str, str] adapter_identity: str | None = None - resume_supported: bool = True - same_world_size_resume: bool | None = True + resume_supported: bool = False torch_runtime_api_version: int = 1 def __post_init__(self) -> None: @@ -2137,17 +1192,9 @@ def __post_init__(self) -> None: _digest_value(digest, f"payload_files[{name}]") if not isinstance(self.resume_supported, bool): raise AlgorithmConfigurationError("resume_supported must be boolean") - if self.same_world_size_resume is not None and not isinstance( - self.same_world_size_resume, bool - ): - raise AlgorithmConfigurationError("same_world_size_resume must be boolean") - if self.resume_supported and self.same_world_size_resume is not True: - raise AlgorithmConfigurationError( - "supported recovery requires same world size" - ) - if not self.resume_supported and self.same_world_size_resume is not None: + if self.resume_supported: raise AlgorithmConfigurationError( - "unsupported recovery must omit same_world_size_resume" + "Torch Runtime API v1 does not support cross-Run recovery" ) object.__setattr__( self, "payload_files", MappingProxyType(dict(sorted(files.items()))) @@ -2174,15 +1221,13 @@ def to_dict(self) -> dict[str, Any]: "resume_supported": self.resume_supported, "torch_runtime_api_version": self.torch_runtime_api_version, } - if self.same_world_size_resume is not None: - payload["same_world_size_resume"] = self.same_world_size_resume return payload @classmethod def from_dict(cls, value: Mapping[str, Any]) -> "TorchCheckpointDescriptor": try: identity = TorchStageRunIdentity.from_dict(value["identity"]) - resume_supported = value.get("resume_supported", True) + resume_supported = value.get("resume_supported", False) return cls( schema_version=value["schema_version"], identity=identity, @@ -2197,9 +1242,6 @@ def from_dict(cls, value: Mapping[str, Any]) -> "TorchCheckpointDescriptor": payload_files=value["payload_files"], adapter_identity=value.get("adapter_identity"), resume_supported=resume_supported, - same_world_size_resume=value.get( - "same_world_size_resume", True if resume_supported else None - ), torch_runtime_api_version=value.get("torch_runtime_api_version", 1), ) except KeyError as exc: @@ -2214,10 +1256,6 @@ def from_dict(cls, value: Mapping[str, Any]) -> "TorchCheckpointDescriptor": "TorchBackwardResult", "TorchCheckpointPayloadDraft", "TorchCheckpointDescriptor", - "TorchCheckpointLocator", - "TorchRecoveryEnvelope", - "TorchCheckpointProgress", - "TorchRankProgressStatistics", "TorchCheckpointRef", "TorchCompositeGlobalState", "TorchCompositeLossContribution", @@ -2229,18 +1267,10 @@ def from_dict(cls, value: Mapping[str, Any]) -> "TorchCheckpointDescriptor": "TorchMetricPolicy", "TorchMetricReductionContext", "TorchMetricReductionResult", - "TorchPreflightLease", - "TorchPreflightTokenData", "TorchStageRunIdentity", "TorchStepLoss", - "TorchRuntimeExecutionEnvelope", - "TorchWorkerControlEnvelope", "apply_torch_loss_backward", - "claim_torch_run_directory", - "describe_torch_checkpoint", "invoke_torch_global_loss_reducer", "reduce_torch_metrics", "report_torch_checkpoint", - "torch_run_config_name", - "validate_torch_retry_identity", ] diff --git a/src/tributo/algorithms/conformance.py b/src/tributo/algorithms/conformance.py index a14d24f..87259fe 100644 --- a/src/tributo/algorithms/conformance.py +++ b/src/tributo/algorithms/conformance.py @@ -3,8 +3,13 @@ from __future__ import annotations import importlib.metadata +import json import sys +from collections.abc import Mapping, Sequence from dataclasses import dataclass, replace +from pathlib import Path + +from packaging.utils import canonicalize_name from tributo.algorithms.api import DistributedAlgorithmDescriptor from tributo.algorithms.core.contracts import validate_contract_binding @@ -102,6 +107,103 @@ def validate_installed_algorithm_package( ) +def _load_identity_manifest(path: Path) -> Mapping[str, Mapping[str, str]]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError("algorithm identity manifest is unavailable") from exc + if not isinstance(payload, Mapping): + raise ValueError("algorithm identity manifest is malformed") + entries = payload.get("entry_points") + if payload.get("schema_version") != 1 or not isinstance(entries, Mapping): + raise ValueError("algorithm identity manifest is malformed") + normalized: dict[str, Mapping[str, str]] = {} + for name, identity in entries.items(): + if not isinstance(name, str) or not isinstance(identity, Mapping): + raise ValueError("algorithm identity manifest entry is malformed") + required = {"distribution", "algorithm_id", "implementation_id"} + if set(identity) != required or any( + not isinstance(identity[field], str) or not identity[field] + for field in required + ): + raise ValueError("algorithm identity manifest entry is incomplete") + normalized[name] = {field: identity[field] for field in sorted(required)} + return normalized + + +def _distribution_name(entry_point: importlib.metadata.EntryPoint) -> str: + distribution = getattr(entry_point, "dist", None) + if distribution is None: + raise ValueError("algorithm Entry Point has no owning distribution") + name = distribution.metadata.get("Name") + if not isinstance(name, str) or not name: + raise ValueError("algorithm Entry Point distribution has no name") + return canonicalize_name(name) + + +def _run_installed_conformance( + *, + distribution_prefix: str, + expected_count: int, + identity_manifest: Path, + required_contracts: Sequence[str], + forbidden_imports: Sequence[str], +) -> tuple[AlgorithmPackageConformanceReport, ...]: + """Validate the installed algorithm Wheels without repository test imports.""" + if not distribution_prefix or expected_count < 1: + raise ValueError("distribution prefix and positive expected count are required") + manifest = _load_identity_manifest(identity_manifest) + entry_points = tuple( + entry_point + for entry_point in importlib.metadata.entry_points(group="tributo.algorithms") + if _distribution_name(entry_point).startswith( + canonicalize_name(distribution_prefix) + ) + ) + if len(entry_points) != expected_count or set(manifest) != { + entry_point.name for entry_point in entry_points + }: + raise ValueError("installed algorithm Entry Point set does not match manifest") + required = tuple(required_contracts) + if tuple(dict.fromkeys(required)) != required or any( + name not in {"config", "input", "output", "coverage"} for name in required + ): + raise ValueError("required contract kinds are invalid") + reports: list[AlgorithmPackageConformanceReport] = [] + for entry_point in sorted(entry_points, key=lambda item: item.name): + descriptor = entry_point.load() + report = validate_installed_algorithm_package( + descriptor, + entry_point_name=entry_point.name, + ) + bindings = descriptor.registration.contract_bindings + if bindings is None or any( + getattr(bindings, name) is None for name in required + ): + raise ValueError( + f"algorithm {entry_point.name!r} is missing required contracts" + ) + actual_identity = { + "algorithm_id": report.algorithm_id, + "distribution": canonicalize_name(report.distribution), + "implementation_id": report.implementation_id, + } + if actual_identity != dict(manifest[entry_point.name]): + raise ValueError( + f"algorithm {entry_point.name!r} identity does not match manifest" + ) + reports.append(report) + imported = sorted( + name + for name in forbidden_imports + if name in sys.modules + or any(module.startswith(f"{name}.") for module in sys.modules) + ) + if imported: + raise ValueError(f"descriptor discovery imported forbidden modules: {imported}") + return tuple(reports) + + __all__ = [ "AlgorithmPackageConformanceReport", "validate_algorithm_descriptor_conformance", diff --git a/src/tributo/algorithms/conformance_cli.py b/src/tributo/algorithms/conformance_cli.py new file mode 100644 index 0000000..1fc54ff --- /dev/null +++ b/src/tributo/algorithms/conformance_cli.py @@ -0,0 +1,48 @@ +"""Source-free CLI for installed algorithm Wheel conformance.""" + +from __future__ import annotations + +import argparse +import json +from collections.abc import Sequence +from dataclasses import asdict +from pathlib import Path + +from tributo.algorithms.conformance import _run_installed_conformance + + +def _comma_separated(value: str) -> tuple[str, ...]: + return tuple(item.strip() for item in value.split(",") if item.strip()) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--distribution-prefix", required=True) + parser.add_argument("--expected-count", required=True, type=int) + parser.add_argument("--identity-manifest", required=True, type=Path) + parser.add_argument("--require-contracts", default="config,input,output,coverage") + parser.add_argument("--forbid-import", default="") + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Run installed-Wheel conformance and print one deterministic report.""" + args = _parser().parse_args(argv) + reports = _run_installed_conformance( + distribution_prefix=args.distribution_prefix, + expected_count=args.expected_count, + identity_manifest=args.identity_manifest, + required_contracts=_comma_separated(args.require_contracts), + forbidden_imports=_comma_separated(args.forbid_import), + ) + print( + json.dumps( + {"count": len(reports), "reports": [asdict(item) for item in reports]}, + sort_keys=True, + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/tributo/algorithms/core/builder.py b/src/tributo/algorithms/core/builder.py index 3da417b..3127d5e 100644 --- a/src/tributo/algorithms/core/builder.py +++ b/src/tributo/algorithms/core/builder.py @@ -305,23 +305,22 @@ def from_torch( execution_plan=SingleStageTorchPlan( stage=TorchStageSpec( stage_id="train", - worker_loop_ref=( - "tributo.integrations.algorithm_runtimes.ray_train_torch:" - "torch_recipe_train_loop_per_worker" - ), input_roles=("train",), ) ), state_layout="replicated", metric_reducers=normalized_reducers, backend=backend, - resume_supported=True, - same_world_size_resume=True, + resume_supported=False, ) if resolved_policy.backend != backend and policy is not None: raise AlgorithmConfigurationError( "Torch Policy backend conflicts with the requested backend" ) + if dict(resolved_policy.metric_reducers) != normalized_reducers: + raise AlgorithmConfigurationError( + "Torch Policy metric reducers conflict with the Builder declaration" + ) return AlgorithmBuilder.from_distributed_algorithm( spec=spec, implementation_id=implementation_id, @@ -402,10 +401,6 @@ def from_torch_adapter( raise AlgorithmConfigurationError( "Torch adapter registrations require policy.loop_owner='adapter'" ) - if not policy.resume_supported and policy.same_world_size_resume is not None: - raise AlgorithmConfigurationError( - "adapter registrations must omit same_world_size_resume when recovery is disabled" - ) return AlgorithmBuilder.from_torch( spec=spec, implementation_id=implementation_id, diff --git a/src/tributo/algorithms/core/dispatcher.py b/src/tributo/algorithms/core/dispatcher.py index 636f452..d21462f 100644 --- a/src/tributo/algorithms/core/dispatcher.py +++ b/src/tributo/algorithms/core/dispatcher.py @@ -22,8 +22,6 @@ ResolvedAlgorithmPlan, StateCoordinationEvidence, TorchExecutionEvidence, - TorchPreflightLease, - TorchRuntimeExecutionEnvelope, WorkerExecutionEvidence, WorkerExecutionResult, WorkerResources, @@ -40,7 +38,6 @@ ResolvedInputLease, RuntimeExecutionEnvelope, RuntimeInputBinding, - TorchRuntimePreflight, WorkerInputPayload, WorkerInputPayloadSet, ) @@ -71,19 +68,10 @@ def execute( artifacts: tuple[ArtifactDraft, ...] = (), cancelled: bool = False, run_id: str | None = None, - torch_preflight_lease: TorchPreflightLease | None = None, ) -> AlgorithmRunResult: """Open input, invoke the selected Runtime, and close in reverse order.""" plan.validate_integrity() run_id = run_id or uuid.uuid4().hex - is_torch = ( - plan.distribution_spec is not None - and plan.distribution_spec.strategy is DistributionStrategy.RAY_TRAIN_TORCH - ) - if is_torch and torch_preflight_lease is None: - raise AlgorithmConfigurationError( - "Torch execution requires a preflight lease" - ) try: runtime = self._runtimes[plan.runtime.runtime_id] except KeyError as exc: @@ -122,14 +110,7 @@ def execute( artifacts=artifacts, cancelled=cancelled, ) - runtime_result = cast(Any, runtime).execute( - TorchRuntimeExecutionEnvelope( - base=base_envelope, - preflight_lease=cast(TorchPreflightLease, torch_preflight_lease), - ) - if is_torch - else base_envelope - ) + runtime_result = cast(Any, runtime).execute(base_envelope) if not isinstance(runtime_result, WorkerExecutionResult) or not isinstance( runtime_result.execution, AlgorithmExecutionResult ): @@ -183,13 +164,6 @@ def execute( lease.close() except Exception as exc: cleanup_errors.append(exc) - if is_torch and torch_preflight_lease is not None: - try: - if torch_preflight_lease.state != "consumed": - torch_preflight_lease.close() - except Exception as exc: - cleanup_errors.append(exc) - if primary_error is not None: for cleanup_error in cleanup_errors: primary_error.add_note( @@ -248,8 +222,7 @@ def torch_preflight( plan: ResolvedAlgorithmPlan, *, run_id: str, - invocation_id: str, - ) -> TorchPreflightLease: + ) -> None: """Run the Torch-only environment check before any input or Ray lease.""" if ( plan.distribution_spec is None @@ -265,25 +238,10 @@ def torch_preflight( raise AlgorithmConfigurationError( f"missing execution component for {exc.args[0]!r}" ) from exc - if not isinstance(runtime, TorchRuntimePreflight): + preflight = getattr(runtime, "preflight", None) + if not callable(preflight): raise AlgorithmConfigurationError("selected Torch runtime has no preflight") - return runtime.preflight(plan, run_id, invocation_id) - - @staticmethod - def claim_torch_preflight( - plan: ResolvedAlgorithmPlan, - *, - run_id: str, - invocation_id: str, - lease: TorchPreflightLease, - ) -> None: - """Claim a preflight lease before opening Runtime or inputs.""" - lease.claim( - run_id=run_id, - invocation_id=invocation_id, - plan_digest=plan.plan_id, - runtime_id=plan.runtime.runtime_id, - ) + preflight(plan, run_id) @staticmethod def _combine_role_bindings( @@ -570,39 +528,17 @@ def execute_plan( if is_torch and cancelled: raise AlgorithmExecutionError("Torch execution was cancelled") run_id = uuid.uuid4().hex if is_torch else None - torch_lease: TorchPreflightLease | None = None if is_torch: torch_run_id = cast(str, run_id) - invocation_id = uuid.uuid4().hex - torch_lease = self._coordinator.torch_preflight( + self._coordinator.torch_preflight(plan, run_id=torch_run_id) + if plan.runtime.execution_profile is None: + return self._coordinator.execute( plan, - run_id=torch_run_id, - invocation_id=invocation_id, + context, + artifacts, + cancelled=cancelled, + run_id=run_id, ) - try: - self._coordinator.claim_torch_preflight( - plan, - run_id=torch_run_id, - invocation_id=invocation_id, - lease=torch_lease, - ) - except BaseException: - torch_lease.close() - raise - if plan.runtime.execution_profile is None: - try: - return self._coordinator.execute( - plan, - context, - artifacts, - cancelled=cancelled, - run_id=run_id, - torch_preflight_lease=torch_lease, - ) - except BaseException: - if torch_lease is not None and torch_lease.state != "consumed": - torch_lease.close() - raise if plan.distribution_spec is None: raise AlgorithmConfigurationError( "formal execution profile requires a DistributionSpec" @@ -613,34 +549,28 @@ def execute_plan( memory_bytes=getattr(plan.runtime, "memory_bytes", None), custom=plan.runtime.custom_resources, ) - try: - with self._runtime_manager.open( - plan.runtime.execution_profile, - resources_per_worker=resources, - worker_count=plan.runtime.worker_count, - ) as runtime_session: - result = self._coordinator.execute( - plan, - context, - artifacts, - cancelled=cancelled, - run_id=run_id, - torch_preflight_lease=torch_lease, - ) - receipt = result.execution_receipt - if receipt is None: - return result - updated_receipt = replace( - cast(ExecutionReceipt, receipt), - cluster_resources=dict(runtime_session.cluster_resources), - runtime_owned=runtime_session.runtime_owned, - resource_preflight=runtime_session.resource_preflight, - ) - return replace(result, execution_receipt=updated_receipt) - except BaseException: - if torch_lease is not None and torch_lease.state != "consumed": - torch_lease.close() - raise + with self._runtime_manager.open( + plan.runtime.execution_profile, + resources_per_worker=resources, + worker_count=plan.runtime.worker_count, + ) as runtime_session: + result = self._coordinator.execute( + plan, + context, + artifacts, + cancelled=cancelled, + run_id=run_id, + ) + receipt = result.execution_receipt + if receipt is None: + return result + updated_receipt = replace( + cast(ExecutionReceipt, receipt), + cluster_resources=dict(runtime_session.cluster_resources), + runtime_owned=runtime_session.runtime_owned, + resource_preflight=runtime_session.resource_preflight, + ) + return replace(result, execution_receipt=updated_receipt) __all__ = ["AlgorithmDispatcher", "AlgorithmRunCoordinator"] diff --git a/src/tributo/algorithms/core/planner.py b/src/tributo/algorithms/core/planner.py index 250fbea..749236a 100644 --- a/src/tributo/algorithms/core/planner.py +++ b/src/tributo/algorithms/core/planner.py @@ -316,11 +316,6 @@ def _resolve_runtime( strategy=distribution.strategy, distribution_digest=distribution.digest, resume_from=request.resume_from, - torch_recovery=( - request.torch_recovery.to_dict() - if request.torch_recovery is not None - else None - ), ) @staticmethod diff --git a/src/tributo/algorithms/spi/__init__.py b/src/tributo/algorithms/spi/__init__.py index 5c76611..9f793c0 100644 --- a/src/tributo/algorithms/spi/__init__.py +++ b/src/tributo/algorithms/spi/__init__.py @@ -16,7 +16,6 @@ PortableRuntimeAdapter, Predictable, RuntimeExecutionEnvelope, - TorchRuntimePreflight, Transformable, ) from tributo.algorithms.spi.input import ( @@ -74,7 +73,6 @@ "PortableRuntimeAdapter", "Predictable", "RuntimeExecutionEnvelope", - "TorchRuntimePreflight", "PreparedInput", "ResolvedInputLease", "RuntimeInputBinding", diff --git a/src/tributo/algorithms/spi/execution.py b/src/tributo/algorithms/spi/execution.py index defbbdd..b59e92d 100644 --- a/src/tributo/algorithms/spi/execution.py +++ b/src/tributo/algorithms/spi/execution.py @@ -16,7 +16,6 @@ WorkerExecutionResult, ) from tributo.algorithms.api.distribution import StateField -from tributo.algorithms.api.torch_runtime import TorchPreflightLease from tributo.algorithms.spi.input import WorkerInputPayload, WorkerInputPayloadSet from tributo.util.annotations import PublicAPI @@ -135,22 +134,6 @@ def runtime_id(self) -> str: ... def execute(self, envelope: RuntimeExecutionEnvelope) -> WorkerExecutionResult: ... -@PublicAPI(stability="alpha") -@runtime_checkable -class TorchRuntimePreflight(Protocol): - """Torch-only preflight surface kept out of the generic Runtime SPI.""" - - @property - def runtime_id(self) -> str: ... - - def preflight( - self, - plan: ResolvedAlgorithmPlan, - run_id: str, - invocation_id: str, - ) -> TorchPreflightLease: ... - - @PublicAPI(stability="alpha") class CollectiveAlgorithm(ABC): """Required surface for iterative Ray Train collective algorithms.""" @@ -439,6 +422,5 @@ def retry_safe(self) -> bool: "PortableRuntimeAdapter", "Predictable", "RuntimeExecutionEnvelope", - "TorchRuntimePreflight", "Transformable", ] diff --git a/src/tributo/algorithms/spi/torch.py b/src/tributo/algorithms/spi/torch.py index 76653b5..c68be0c 100644 --- a/src/tributo/algorithms/spi/torch.py +++ b/src/tributo/algorithms/spi/torch.py @@ -2,6 +2,7 @@ from __future__ import annotations +import math from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import dataclass, field @@ -34,12 +35,10 @@ class TorchRuntimeContext: input_bindings: Mapping[str, object] = field(default_factory=dict) output_config: Mapping[str, object] = field(default_factory=dict) - recovery_identity: Mapping[str, object] | None = None input_binding_digest: str | None = None state_layout: str = "replicated" adapter_identity: str | None = None - resume_supported: bool = True - same_world_size_resume: bool | None = True + resume_supported: bool = False torch_runtime_api_version: int = 1 def __post_init__(self) -> None: @@ -75,21 +74,11 @@ def __post_init__(self) -> None: raise ValueError("TorchRuntimeContext state_layout is invalid") if not isinstance(self.resume_supported, bool): raise ValueError("TorchRuntimeContext resume_supported must be boolean") - if self.same_world_size_resume is not None and not isinstance( - self.same_world_size_resume, bool - ): - raise ValueError( - "TorchRuntimeContext same_world_size_resume must be boolean" - ) + if self.resume_supported: + raise ValueError("Torch Runtime API v1 does not support cross-Run recovery") object.__setattr__(self, "algorithm_config", deep_freeze(self.algorithm_config)) object.__setattr__(self, "input_bindings", deep_freeze(self.input_bindings)) object.__setattr__(self, "output_config", deep_freeze(self.output_config)) - if self.recovery_identity is not None: - object.__setattr__( - self, - "recovery_identity", - deep_freeze(self.recovery_identity), - ) def to_dict(self) -> dict[str, object]: payload: dict[str, object] = { @@ -101,17 +90,12 @@ def to_dict(self) -> dict[str, object]: "run_identity": self.run_identity.to_dict() if self.run_identity else None, "input_bindings": dict(self.input_bindings), "output_config": dict(self.output_config), - "recovery_identity": dict(self.recovery_identity) - if self.recovery_identity - else None, "input_binding_digest": self.input_binding_digest, "state_layout": self.state_layout, "adapter_identity": self.adapter_identity, "resume_supported": self.resume_supported, "torch_runtime_api_version": self.torch_runtime_api_version, } - if self.same_world_size_resume is not None: - payload["same_world_size_resume"] = self.same_world_size_resume return payload @@ -127,9 +111,6 @@ class TorchStageContext: input_roles: tuple[str, ...] predecessor_stage_id: str | None = None predecessor_checkpoint_descriptor: Mapping[str, Any] | None = None - metric_mapping: Mapping[str, str] = field(default_factory=dict) - checkpoint_required: bool = True - checkpoint_interval_windows: int = 1 def __post_init__(self) -> None: if not isinstance(self.runtime, TorchRuntimeContext): @@ -145,27 +126,6 @@ def __post_init__(self) -> None: if not isinstance(self.is_final, bool): raise ValueError("TorchStageContext is_final must be boolean") object.__setattr__(self, "input_roles", tuple(self.input_roles)) - if any( - not isinstance(name, str) - or not name - or not isinstance(target, str) - or not target - for name, target in self.metric_mapping.items() - ): - raise ValueError("TorchStageContext metric_mapping is malformed") - if len(set(self.metric_mapping.values())) != len(self.metric_mapping): - raise ValueError("TorchStageContext metric_mapping targets must be unique") - if not isinstance(self.checkpoint_required, bool): - raise ValueError("TorchStageContext checkpoint_required must be boolean") - if ( - not isinstance(self.checkpoint_interval_windows, int) - or isinstance(self.checkpoint_interval_windows, bool) - or self.checkpoint_interval_windows < 1 - ): - raise ValueError( - "TorchStageContext checkpoint_interval_windows must be positive" - ) - object.__setattr__(self, "metric_mapping", deep_freeze(self.metric_mapping)) if self.predecessor_checkpoint_descriptor is not None: if any( key in {"locator", "checkpoint_locator", "path", "credential"} @@ -193,9 +153,6 @@ def to_dict(self) -> dict[str, object]: ) if self.predecessor_checkpoint_descriptor else None, - "metric_mapping": dict(self.metric_mapping), - "checkpoint_required": self.checkpoint_required, - "checkpoint_interval_windows": self.checkpoint_interval_windows, } @classmethod @@ -204,7 +161,7 @@ def from_dict(cls, value: Mapping[str, object]) -> "TorchStageContext": if not isinstance(runtime_value, Mapping): raise ValueError("TorchStageContext runtime payload is invalid") identity_value = runtime_value.get("run_identity") - resume_supported = runtime_value.get("resume_supported", True) + resume_supported = runtime_value.get("resume_supported", False) runtime = TorchRuntimeContext( algorithm_config=runtime_value.get("algorithm_config", {}), implementation_id=runtime_value["implementation_id"], @@ -218,14 +175,10 @@ def from_dict(cls, value: Mapping[str, object]) -> "TorchStageContext": ), input_bindings=runtime_value.get("input_bindings", {}), output_config=runtime_value.get("output_config", {}), - recovery_identity=runtime_value.get("recovery_identity"), input_binding_digest=runtime_value.get("input_binding_digest"), state_layout=runtime_value.get("state_layout", "replicated"), adapter_identity=runtime_value.get("adapter_identity"), resume_supported=resume_supported, - same_world_size_resume=runtime_value.get( - "same_world_size_resume", True if resume_supported else None - ), torch_runtime_api_version=runtime_value.get("torch_runtime_api_version", 1), ) return cls( @@ -239,11 +192,6 @@ def from_dict(cls, value: Mapping[str, object]) -> "TorchStageContext": Mapping[str, Any] | None, value.get("predecessor_checkpoint_descriptor"), ), - metric_mapping=cast(Mapping[str, str], value.get("metric_mapping", {})), - checkpoint_required=cast(bool, value.get("checkpoint_required", True)), - checkpoint_interval_windows=cast( - int, value.get("checkpoint_interval_windows", 1) - ), ) @@ -279,9 +227,7 @@ def __post_init__(self) -> None: raise ValueError("TorchWorkerCheckpointContext stage is required") if self.source not in { "none", - "ray_failure_retry", "stage_dependency", - "cross_run_initial_recovery", }: raise ValueError("TorchWorkerCheckpointContext source is invalid") if self.source == "none" and self.checkpoint is not None: @@ -514,6 +460,13 @@ def __post_init__(self) -> None: or self.gradient_accumulation_steps < 1 ): raise ValueError("gradient_accumulation_steps must be positive") + if self.max_gradient_norm is not None and ( + not isinstance(self.max_gradient_norm, (int, float)) + or isinstance(self.max_gradient_norm, bool) + or not math.isfinite(float(self.max_gradient_norm)) + or float(self.max_gradient_norm) <= 0 + ): + raise ValueError("max_gradient_norm must be finite and positive") @PublicAPI(stability="alpha") diff --git a/src/tributo/inference/kernel.py b/src/tributo/inference/kernel.py index edd4f6e..3c7c531 100644 --- a/src/tributo/inference/kernel.py +++ b/src/tributo/inference/kernel.py @@ -153,14 +153,6 @@ def _build_input_tensor( "scalar single-column input must be a one-dimensional batch column" ) tensor = arrays[0] - elif len(arrays) == 1 and arrays[0].dtype == object: - # Arrow/Parquet may decode a vector-valued column as either a one- or - # multi-dimensional object array whose rows contain nested arrays. - # Stack rows before dtype conversion so the declared tensor rank is - # preserved instead of treating the object dimension as a feature axis. - tensor = _stack_nested_object_array(arrays[0]) - if tensor.dtype == object: - tensor = np.column_stack(arrays) elif len(arrays) == 1 and arrays[0].ndim > 1: tensor = arrays[0] else: @@ -181,20 +173,6 @@ def _build_input_tensor( return np.asarray(tensor) -def _stack_nested_object_array(value: object) -> np.ndarray: - """Materialize homogeneous nested object arrays while preserving row rank.""" - array = np.asarray(value) - if array.dtype != object: - return array - try: - nested = [_stack_nested_object_array(item) for item in array.tolist()] - if nested: - return np.stack(nested) - except (TypeError, ValueError): - pass - return array - - def _tensor_row_count(tensors: dict[str, np.ndarray], *, kind: str) -> int: row_count: int | None = None for name, tensor in tensors.items(): diff --git a/src/tributo/integrations/algorithm_runtimes/ray_train_torch.py b/src/tributo/integrations/algorithm_runtimes/ray_train_torch.py index b452a5b..b8c319c 100644 --- a/src/tributo/integrations/algorithm_runtimes/ray_train_torch.py +++ b/src/tributo/integrations/algorithm_runtimes/ray_train_torch.py @@ -12,12 +12,12 @@ import json import logging import math -import os import tempfile from collections.abc import Mapping -from contextlib import contextmanager, nullcontext +from contextlib import nullcontext +from dataclasses import replace from pathlib import Path -from typing import Any, Generator, cast +from typing import Any, cast from tributo.algorithms.api import ( AlgorithmConfigurationError, @@ -31,8 +31,7 @@ TorchAccumulationWindow, TorchBackwardContext, TorchCheckpointDescriptor, - TorchCheckpointLocator, - TorchCheckpointProgress, + TorchCheckpointPayloadDraft, TorchCheckpointRef, TorchCompositeGlobalState, TorchCompositeLossContribution, @@ -42,25 +41,16 @@ TorchGlobalLossReduction, TorchLossContribution, TorchMetricContribution, - TorchPreflightLease, - TorchPreflightTokenData, - TorchRankProgressStatistics, - TorchRecoveryEnvelope, TorchRoleExecutionEvidence, - TorchRuntimeExecutionEnvelope, TorchStageRunIdentity, - TorchWorkerControlEnvelope, WorkerExecutionEvidence, WorkerExecutionResult, apply_torch_loss_backward, - claim_torch_run_directory, - describe_torch_checkpoint, invoke_torch_global_loss_reducer, report_torch_checkpoint, - torch_run_config_name, - validate_torch_retry_identity, ) from tributo.algorithms.api.models import FORMAL_DISTRIBUTED_STRATEGY_CONTRACTS +from tributo.algorithms.api.torch_runtime import _describe_torch_checkpoint from tributo.algorithms.core.worker import ( _actual_environment_versions, _load_reference, @@ -92,14 +82,6 @@ RAY_TRAIN_TORCH_RUNTIME_ID = FORMAL_DISTRIBUTED_STRATEGY_CONTRACTS[ DistributionStrategy.RAY_TRAIN_TORCH ].runtime_id -_CORE_RECIPE_LOOP_REF = ( - "tributo.integrations.algorithm_runtimes.ray_train_torch:" - "torch_recipe_train_loop_per_worker" -) -_CORE_ADAPTER_LOOP_REF = ( - "tributo.integrations.algorithm_runtimes.ray_train_torch:" - "ray_torch_adapter_train_loop_per_worker" -) logger = logging.getLogger(__name__) @@ -146,6 +128,29 @@ def _policy(plan: Any) -> Any: return policy +def _validate_runtime_policy(policy: Any) -> None: + """Reject reserved Torch Policy features not implemented by Runtime v1.""" + if policy.checkpoint_owner_rank != 0: + raise AlgorithmConfigurationError( + "Torch Runtime API v1 requires checkpoint_owner_rank=0" + ) + if policy.state_layout == "sharded": + raise AlgorithmConfigurationError( + "Torch sharded state is reserved and not supported by Runtime v1" + ) + if any(route.mode == "split_framework" for route in policy.dataset_routing): + raise AlgorithmConfigurationError( + "Torch split_framework routing is reserved and not supported by Runtime v1" + ) + if ( + policy.checkpoint_adapter_ref is not None + or policy.evidence_adapter_ref is not None + ): + raise AlgorithmConfigurationError( + "Torch checkpoint/evidence adapter references are reserved until their protocols are gated" + ) + + def _torch_algorithm_context_config(plan: Any) -> dict[str, object]: """Return only algorithm-owned config for Torch implementation contexts.""" config = plan.algorithm_config @@ -175,6 +180,35 @@ def _torch_output_config(plan: Any) -> dict[str, object]: return {str(key): value for key, value in output.items()} +def _torch_ray_config(plan: Any) -> tuple[str | None, int]: + """Validate and return only Ray options consumed by Runtime API v1.""" + value = plan.algorithm_config.get("ray", {}) + if not isinstance(value, Mapping): + raise AlgorithmConfigurationError("Torch ray config must be a mapping") + unknown = sorted(set(value) - {"storage_path", "max_failures"}) + if unknown: + raise AlgorithmConfigurationError( + f"Torch ray config contains unsupported key(s): {unknown}" + ) + storage_path = value.get("storage_path") + if storage_path is not None and ( + not isinstance(storage_path, str) or not storage_path + ): + raise AlgorithmConfigurationError( + "Torch ray.storage_path must be a non-empty string" + ) + max_failures = value.get("max_failures", 0) + if ( + not isinstance(max_failures, int) + or isinstance(max_failures, bool) + or max_failures < -1 + ): + raise AlgorithmConfigurationError( + "ray.max_failures must be -1 or a non-negative integer" + ) + return storage_path, max_failures + + _ADAPTER_CONFIG_BLOCKED_KEYS = frozenset( { "ray", @@ -252,6 +286,31 @@ def _accumulate_metric_totals( totals[1] += contribution.normalizer +def _record_stage_metrics( + target: dict[str, float], + metrics: Mapping[str, object], + declared_names: set[str], + *, + is_final: bool, +) -> None: + """Retain directly named Policy metrics across a component plan.""" + for name in declared_names: + if name == "train_loss" and not is_final: + continue + value = metrics.get(name) + if value is None: + continue + if ( + not isinstance(value, (int, float)) + or isinstance(value, bool) + or not math.isfinite(float(value)) + ): + raise AlgorithmExecutionError( + f"Torch metric {name!r} must be a finite number" + ) + target[name] = float(value) + + def _reduce_metric_totals( totals: Mapping[str, list[float]], reducers: Mapping[str, str], @@ -273,10 +332,6 @@ def _reduce_metric_totals( for name in sorted(names): numerator, normalizer = totals[name] reducer_value = reducers.get(name) - if reducer_value is None and "_" in name: - reducer_value = reducers.get(name.split("_", 1)[1]) - if reducer_value is None and name.endswith("_loss"): - reducer_value = reducers.get("train_loss") if reducer_value is None: raise AlgorithmConfigurationError( f"Torch metric {name!r} has no declared reducer" @@ -397,287 +452,15 @@ def _stage_context( input_roles=tuple(stage.input_roles), predecessor_stage_id=predecessor, predecessor_checkpoint_descriptor=predecessor_descriptor, - metric_mapping=dict(getattr(stage, "metric_mapping", {})), - checkpoint_required=bool(getattr(stage, "checkpoint_required", True)), - checkpoint_interval_windows=int( - getattr(stage, "checkpoint_interval_windows", 1) - ), ) -def _control_for_stage( - plan: Any, - policy: Any, - stage: Any, - *, - run_id: str, - invocation_id: str, - checkpoint: Mapping[str, Any] | None = None, - purpose: str | None = None, - source_stage_id: str | None = None, - predecessor: Mapping[str, Any] | None = None, -) -> dict[str, Any] | None: - """Create credential-free initial recovery control for a Stage.""" - if checkpoint is None: - checkpoint = predecessor - if checkpoint is not None and purpose is None: - purpose = "stage_dependency" - if checkpoint is not None and source_stage_id is None: - source_stage_id = getattr(stage, "checkpoint_from_stage", None) - if checkpoint is None: - return None - resume_uri = checkpoint.get("locator") - descriptor_digest = checkpoint.get("descriptor_digest") - if not isinstance(resume_uri, str) or not resume_uri: - raise AlgorithmConfigurationError( - "Torch recovery requires a credential-free locator" - ) - if not isinstance(descriptor_digest, str): - raise AlgorithmConfigurationError( - "Torch recovery requires checkpoint_descriptor_digest" - ) - locator = TorchCheckpointLocator(resume_uri, descriptor_digest) - control = TorchWorkerControlEnvelope( - schema_version=1, - run_id=run_id, - invocation_id=invocation_id, - source_stage_id=source_stage_id, - target_stage_id=stage.stage_id, - purpose=purpose or "cross_run_initial_recovery", - checkpoint_locator=locator, - checkpoint_descriptor_digest=descriptor_digest, - policy_digest=policy.digest, - execution_plan_digest=policy.execution_plan.digest, - ) - return control.to_dict() - - -def _describe_recovery_locator( - locator: TorchCheckpointLocator, - *, - policy: Any, - plan: Any, - worker_count: int, -) -> TorchCheckpointDescriptor: - """Open a recovery locator on the driver and validate its payload digest.""" - checkpoint = open_torch_checkpoint_locator(locator) - try: - descriptor = describe_torch_checkpoint(TorchCheckpointRef(checkpoint), object()) - _require_checkpoint_commit(checkpoint, descriptor) - finally: - closer = getattr(checkpoint, "close", None) - if callable(closer): - closer() - if descriptor.digest != locator.descriptor_digest: - raise AlgorithmExecutionError( - "Torch recovery locator descriptor digest drifted" - ) - if ( - descriptor.policy_digest != policy.digest - or descriptor.execution_plan_digest != policy.execution_plan.digest - or descriptor.world_size != worker_count - or descriptor.implementation_code_digest != plan.implementation.code_digest - or descriptor.identity.plan_digest != plan.plan_id - or descriptor.input_binding_digest != _input_binding_digest(plan) - or descriptor.state_layout != policy.state_layout - ): - raise AlgorithmExecutionError("Torch recovery descriptor identity mismatch") - return descriptor - - -def _checkpoint_evidence_payload(checkpoint: object) -> dict[str, Any]: - """Read the Core-owned, credential-free evidence sidecar when present.""" - opener = getattr(checkpoint, "as_directory", None) - if not callable(opener): - return {} - try: - with opener() as directory: - path = Path(directory) / "torch_execution_evidence.json" - if path.is_symlink() or not path.is_file(): - return {} - payload = json.loads(path.read_text(encoding="utf-8")) - except (OSError, TypeError, ValueError) as exc: - raise AlgorithmExecutionError( - "Torch checkpoint execution evidence is malformed" - ) from exc - if not isinstance(payload, Mapping): - raise AlgorithmExecutionError( - "Torch checkpoint execution evidence is malformed" - ) - return dict(payload) - - -def _require_checkpoint_commit( - checkpoint: object, descriptor: TorchCheckpointDescriptor -) -> None: - """Require the marker-last commit used by persistent Stage locators.""" - with _opened_checkpoint(checkpoint) as root: - marker = root / "torch_stage_commit.json" - if marker.is_symlink() or not marker.is_file(): - raise AlgorithmExecutionError( - "Torch Stage locator references an uncommitted checkpoint" - ) - try: - payload = json.loads(marker.read_text(encoding="utf-8")) - except (OSError, TypeError, ValueError) as exc: - raise AlgorithmExecutionError( - "Torch Stage checkpoint commit marker is malformed" - ) from exc - if not isinstance(payload, Mapping) or ( - payload.get("identity") != descriptor.identity.to_dict() - or payload.get("descriptor_digest") != descriptor.digest - ): - raise AlgorithmExecutionError( - "Torch Stage checkpoint commit marker does not match descriptor" - ) - - -def _recovery_record_for_locator( - locator: TorchCheckpointLocator, - *, - policy: Any, - plan: Any, - worker_count: int, -) -> dict[str, Any]: - descriptor = _describe_recovery_locator( - locator, policy=policy, plan=plan, worker_count=worker_count - ) - checkpoint = open_torch_checkpoint_locator(locator) - try: - evidence = _checkpoint_evidence_payload(checkpoint) - finally: - closer = getattr(checkpoint, "close", None) - if callable(closer): - closer() - return { - "locator": locator.uri, - "descriptor_digest": descriptor.digest, - "descriptor": descriptor.to_dict(), - "evidence": evidence, - } - - -def _recovery_records( - plan: Any, - policy: Any, - *, - worker_count: int, -) -> tuple[tuple[str, ...], str | None, dict[str, dict[str, Any]]]: - """Normalize the full Torch recovery envelope and legacy shorthand.""" - stages = tuple(policy.execution_plan.stages) - stage_ids = tuple(stage.stage_id for stage in stages) - raw_envelope = plan.runtime.torch_recovery - if raw_envelope is not None: - envelope = TorchRecoveryEnvelope.from_dict(raw_envelope) - if not policy.resume_supported and ( - envelope.stage_checkpoints or envelope.active_checkpoint is not None - ): - raise AlgorithmConfigurationError( - "Torch Policy does not support external recovery" - ) - completed = tuple(envelope.completed_stage_ids) - if any(stage_id not in stage_ids for stage_id in completed): - raise AlgorithmExecutionError( - "Torch recovery contains an unknown completed Stage" - ) - if tuple(stage_ids[: len(completed)]) != completed: - raise AlgorithmExecutionError( - "Torch recovery completed stages must follow execution plan order" - ) - if ( - envelope.active_stage_id is not None - and envelope.active_stage_id not in stage_ids - ): - raise AlgorithmExecutionError("Torch recovery active Stage is unknown") - if envelope.active_stage_id is not None: - active_index = stage_ids.index(envelope.active_stage_id) - if any(stage_id not in completed for stage_id in stage_ids[:active_index]): - raise AlgorithmExecutionError( - "Torch recovery active Stage has an incomplete predecessor" - ) - active_stage = stages[active_index] - if any(dep not in completed for dep in active_stage.depends_on): - raise AlgorithmExecutionError( - "Torch recovery active Stage dependencies are incomplete" - ) - records: dict[str, dict[str, Any]] = {} - for stage_id, locator in envelope.stage_checkpoints.items(): - record = _recovery_record_for_locator( - locator, policy=policy, plan=plan, worker_count=worker_count - ) - descriptor = TorchCheckpointDescriptor.from_dict(record["descriptor"]) - if descriptor.identity.stage_id != stage_id: - raise AlgorithmExecutionError( - "Torch recovery checkpoint Stage mismatch" - ) - records[stage_id] = record - if ( - envelope.active_stage_id is not None - and envelope.active_checkpoint is not None - ): - active_record = _recovery_record_for_locator( - envelope.active_checkpoint, - policy=policy, - plan=plan, - worker_count=worker_count, - ) - active_descriptor = TorchCheckpointDescriptor.from_dict( - active_record["descriptor"] - ) - if active_descriptor.identity.stage_id != envelope.active_stage_id: - raise AlgorithmExecutionError( - "Torch recovery active checkpoint Stage mismatch" - ) - if not policy.resume_supported: - raise AlgorithmConfigurationError( - "Torch Policy does not support cross-Run active recovery" - ) - if not active_descriptor.resume_supported: - raise AlgorithmExecutionError( - "Torch active recovery checkpoint is not externally recoverable" - ) - records[envelope.active_stage_id] = active_record - return completed, envelope.active_stage_id, records - - resume_uri = plan.runtime.resume_from - ray_config = plan.algorithm_config.get("ray", {}) - resume_config = ( - ray_config.get("resume", {}) if isinstance(ray_config, Mapping) else {} - ) - if resume_uri is None and isinstance(resume_config, Mapping): - for key in ("uri", "checkpoint_uri"): - if isinstance(resume_config.get(key), str): - resume_uri = resume_config[key] - break - if resume_uri is None: - return (), None, {} - if not policy.resume_supported: - raise AlgorithmConfigurationError( - "Torch Policy does not support external recovery" - ) - descriptor_digest = ( - resume_config.get("checkpoint_descriptor_digest") - if isinstance(resume_config, Mapping) - else None - ) - if not isinstance(descriptor_digest, str): - raise AlgorithmConfigurationError( - "Torch resume shorthand requires ray.resume.checkpoint_descriptor_digest" - ) - locator = TorchCheckpointLocator(resume_uri, descriptor_digest) - record = _recovery_record_for_locator( - locator, policy=policy, plan=plan, worker_count=worker_count +def _worker_stage_context(context: TorchStageContext) -> TorchStageContext: + """Remove Driver-only output paths before serializing a Worker context.""" + return replace( + context, + runtime=replace(context.runtime, output_config={}), ) - descriptor = TorchCheckpointDescriptor.from_dict(record["descriptor"]) - if descriptor.identity.stage_id not in stage_ids: - raise AlgorithmExecutionError( - "Torch resume checkpoint Stage is not in execution plan" - ) - if not descriptor.resume_supported: - raise AlgorithmExecutionError( - "Torch resume checkpoint is not externally recoverable" - ) - return (), descriptor.identity.stage_id, {descriptor.identity.stage_id: record} def _payload_rows(value: object) -> int | None: @@ -716,30 +499,16 @@ def _resource_map(plan: Any) -> dict[str, float]: return resources -@contextmanager -def _opened_checkpoint(checkpoint: object) -> Generator[Path, None, None]: - if isinstance(checkpoint, (str, Path)): - root = Path(checkpoint) - if not root.is_dir(): - raise AlgorithmExecutionError("Torch checkpoint directory is missing") - yield root - return - opener = getattr(checkpoint, "as_directory", None) - if not callable(opener): - raise AlgorithmExecutionError("Torch checkpoint cannot be opened") - with opener() as directory: - yield Path(directory) - - def _validate_stage_routes( policy: Any, stage: Any, - datasets: Mapping[str, object], + datasets: dict[str, object], worker_count: int, -) -> dict[str, int]: +) -> tuple[dict[str, int], dict[str, int]]: """Validate role presence, exact coverage minimums and replication budgets.""" routes = {route.role: route for route in policy.dataset_routing} rows: dict[str, int] = {} + replicated_bytes_by_role: dict[str, int] = {} replicated_bytes = 0 for role in stage.input_roles: route = routes.get(role) @@ -752,7 +521,30 @@ def _validate_stage_routes( f"required Torch role {role!r} is absent" ) continue - count = _payload_rows(dataset) + bounded_dataset = dataset + if route.mode == "replicate": + if route.max_rows is None: + raise AlgorithmConfigurationError( + f"Torch replicate role {role!r} has no max_rows" + ) + limiter = getattr(dataset, "limit", None) + if not callable(limiter): + raise AlgorithmConfigurationError( + f"Torch replicate role {role!r} cannot be bounded" + ) + limited = limiter(route.max_rows + 1) + materialize = getattr(limited, "materialize", None) + if not callable(materialize): + raise AlgorithmConfigurationError( + f"Torch replicate role {role!r} cannot be materialized" + ) + try: + bounded_dataset = materialize() + except Exception as exc: + raise AlgorithmConfigurationError( + f"Torch replicate role {role!r} materialization failed" + ) from exc + count = _payload_rows(bounded_dataset) if count is None: raise AlgorithmConfigurationError( f"Torch role {role!r} row count is not verifiable" @@ -773,25 +565,7 @@ def _validate_stage_routes( raise AlgorithmConfigurationError( f"Torch replicate role {role!r} exceeds max_rows" ) - limiter = getattr(dataset, "limit", None) - if not callable(limiter): - raise AlgorithmConfigurationError( - f"Torch replicate role {role!r} cannot be bounded" - ) - probe = limiter(route.max_rows + 1) - probe_count = getattr(probe, "count", None) - if not callable(probe_count): - raise AlgorithmConfigurationError( - f"Torch replicate role {role!r} bounded size is not verifiable" - ) - observed_probe = int(probe_count()) - if observed_probe > route.max_rows: - raise AlgorithmConfigurationError( - f"Torch replicate role {role!r} exceeds max_rows" - ) - if isinstance(datasets, dict): - datasets[role] = limiter(route.max_rows) - size_bytes = getattr(dataset, "size_bytes", None) + size_bytes = getattr(bounded_dataset, "size_bytes", None) if callable(size_bytes): size_bytes = size_bytes() if not isinstance(size_bytes, int) or size_bytes < 0: @@ -805,7 +579,9 @@ def _validate_stage_routes( raise AlgorithmConfigurationError( f"Torch replicate role {role!r} exceeds max_bytes_per_worker" ) + replicated_bytes_by_role[role] = size_bytes replicated_bytes += size_bytes + datasets[role] = bounded_dataset if ( policy.max_replicated_bytes_per_worker is not None and replicated_bytes > policy.max_replicated_bytes_per_worker @@ -813,80 +589,11 @@ def _validate_stage_routes( raise AlgorithmConfigurationError( "Torch replicate roles exceed aggregate byte budget" ) - return rows - - -def _worker_rows(payload: object, role: str) -> int: - if hasattr(payload, "get") and callable(payload.get): - payload = payload.get(role) - value = getattr(payload, "value", payload) - rows = _payload_rows(value) - return rows if rows is not None else 0 - - -def _worker_evidence( - metrics: Mapping[str, Any], - plan: Any, - identity: TorchStageRunIdentity, - stage: Any, - expected_rows: Mapping[str, int], -) -> tuple[dict[str, Any], ...]: - raw_workers = metrics.get("execution_workers") - if ( - not isinstance(raw_workers, (list, tuple)) - or len(raw_workers) != plan.runtime.worker_count - ): - raise AlgorithmExecutionError("Torch execution did not report every worker") - workers = tuple( - WorkerExecutionEvidence.from_dict(item) - for item in _normalize_worker_evidence(raw_workers, plan) - ) - role_evidence: list[TorchRoleExecutionEvidence] = [] - for role in stage.input_roles: - observed = sum(item.input_rows.get(role, 0) for item in workers) - role_evidence.append( - TorchRoleExecutionEvidence( - role=role, - mode="split_exact", - required=True, - present=True, - empty_rank_policy="reject", - expected_rows=expected_rows.get(role), - observed_rows=observed, - rows_per_rank=tuple(item.input_rows.get(role, 0) for item in workers), - binding_digest=_binding_digest_for_role(plan, role), - ) - ) - global_digest = metrics.get("model_state_digest") - if not isinstance(global_digest, str) or len(global_digest) != 64: - raise AlgorithmExecutionError("Torch execution did not report a model digest") - evidence = TorchExecutionEvidence( - identity=identity, - run_config_name=torch_run_config_name(identity), - policy_digest=_policy(plan).digest, - parallelism_id=_policy(plan).parallelism_id, - state_layout=_policy(plan).state_layout, - workers=workers, - roles=tuple(role_evidence), - replicated_state=( - __import__( - "tributo.algorithms.api.execution", - fromlist=["ReplicatedTorchStateEvidence"], - ).ReplicatedTorchStateEvidence( - model_digests_by_rank={ - item.rank: item.model_state_digest for item in workers - }, - global_model_digest=global_digest, - ) - if _policy(plan).state_layout == "replicated" - else None - ), - ) - return (evidence.to_dict(),) + return rows, replicated_bytes_by_role def _binding_digest_for_role(plan: Any, role: str) -> str: - """Resolve a role digest, falling back to the primary binding for aliases.""" + """Resolve the digest of one actually bound input role.""" if hasattr(plan.input_descriptors, "get"): try: descriptor = plan.input_descriptors.get(role) @@ -894,7 +601,9 @@ def _binding_digest_for_role(plan: Any, role: str) -> str: descriptor = None if descriptor is not None: return cast(str, descriptor.binding_digest) - return cast(str, plan.primary_input_descriptor.binding_digest) + raise AlgorithmConfigurationError( + f"Torch input role {role!r} has no binding digest" + ) def _input_binding_digest(plan: Any) -> str: @@ -947,6 +656,7 @@ def _component_stage_evidence( identity: TorchStageRunIdentity, metrics: Mapping[str, Any], expected_rows: Mapping[str, int], + replicated_bytes_by_role: Mapping[str, int], ) -> ComponentStageEvidence: raw_workers = metrics.get("execution_workers") if not isinstance(raw_workers, (list, tuple)): @@ -963,14 +673,13 @@ def _component_stage_evidence( stage=stage, workers=workers, expected_rows=expected_rows, + replicated_bytes_by_role=replicated_bytes_by_role, ) state_digest = metrics.get("model_state_digest") descriptor = metrics.get("checkpoint_descriptor") if not isinstance(state_digest, str) or len(state_digest) != 64: raise AlgorithmExecutionError("Torch Stage did not report a state digest") - if getattr(stage, "checkpoint_required", True) and not isinstance( - descriptor, Mapping - ): + if not isinstance(descriptor, Mapping): raise AlgorithmExecutionError( "Torch Stage did not report a checkpoint descriptor" ) @@ -988,60 +697,16 @@ def _component_stage_evidence( ) -def _recovered_stage_evidence( - *, - plan: Any, - policy: Any, - stage: Any, - descriptor: Mapping[str, Any], - evidence: Mapping[str, Any], -) -> ComponentStageEvidence: - """Rebuild component evidence persisted beside a completed Stage checkpoint.""" - workers_value = evidence.get("execution_workers") - state_digest = evidence.get("model_state_digest") - if not isinstance(workers_value, (list, tuple)) or not isinstance( - state_digest, str - ): - raise AlgorithmExecutionError( - f"Torch recovery checkpoint for Stage {stage.stage_id!r} has no evidence" - ) - workers = tuple( - WorkerExecutionEvidence.from_dict(item) - for item in _normalize_worker_evidence(workers_value, plan) - ) - if len(workers) != plan.runtime.worker_count: - raise AlgorithmExecutionError( - "Torch recovery Stage worker evidence is incomplete" - ) - expected_rows: dict[str, int] = {} - routes = {route.role: route for route in policy.dataset_routing} - for role in stage.input_roles: - route = routes[role] - rows = tuple(worker.input_rows.get(role, 0) for worker in workers) - if route.mode == "replicate": - if rows and len(set(rows)) == 1: - expected_rows[role] = rows[0] - elif sum(rows) > 0: - expected_rows[role] = sum(rows) - metrics = dict(evidence) - metrics["checkpoint_descriptor"] = dict(descriptor) - identity = TorchStageRunIdentity.from_dict(descriptor["identity"]) - return _component_stage_evidence( - plan=plan, - policy=policy, - stage=stage, - identity=identity, - metrics=metrics, - expected_rows=expected_rows, - ) - - def _component_state_details( stages: tuple[ComponentStageEvidence, ...], + final_stage_id: str | None = None, ) -> dict[str, str | int]: """Project component evidence into the scalar Core state receipt details.""" if not stages: raise AlgorithmExecutionError("Torch component state requires Stage evidence") + anchor_stage = final_stage_id or stages[-1].stage_id + if anchor_stage not in {stage.stage_id for stage in stages}: + raise AlgorithmExecutionError("Torch component final Stage evidence is missing") composition_digest = hashlib.sha256( json.dumps( [stage.to_dict() for stage in stages], @@ -1053,7 +718,7 @@ def _component_state_details( "framework": "torch_component", "component_stage_count": len(stages), "component_stages": ",".join(stage.stage_id for stage in stages), - "anchor_stage": stages[-1].stage_id, + "anchor_stage": anchor_stage, "composition_digest": composition_digest, } for stage in stages: @@ -1076,6 +741,7 @@ def _role_execution_evidence( stage: Any, workers: tuple[WorkerExecutionEvidence, ...], expected_rows: Mapping[str, int], + replicated_bytes_by_role: Mapping[str, int], ) -> tuple[TorchRoleExecutionEvidence, ...]: """Build role evidence from the Policy route instead of a default split.""" routes = {route.role: route for route in policy.dataset_routing} @@ -1092,7 +758,7 @@ def _role_execution_evidence( f"Torch replicate role {role!r} is not identical across workers" ) observed_rows = rows_per_rank[0] - replicated_bytes = route.max_bytes_per_worker + replicated_bytes = replicated_bytes_by_role.get(role) else: observed_rows = sum(rows_per_rank) replicated_bytes = None @@ -1115,7 +781,9 @@ def _role_execution_evidence( observed_rows=observed_rows, rows_per_rank=rows_per_rank, replicated_bytes_per_worker=replicated_bytes, - binding_digest=_binding_digest_for_role(plan, role), + binding_digest=( + _binding_digest_for_role(plan, role) if present else None + ), ) ) return tuple(evidence) @@ -1129,7 +797,7 @@ def _reduce_composite_loss( device: object, dist: Any, observation: dict[str, object] | None = None, - expected_metrics: frozenset[str] = frozenset(), + expected_reducer_metrics: frozenset[str] = frozenset({"train_loss"}), ) -> TorchGlobalLossReduction: """AllReduce generic component state, then invoke the Wheel-owned reducer.""" import torch @@ -1232,9 +900,9 @@ def _reduce_composite_loss( raise AlgorithmExecutionError( f"Composite loss reducer rejected contribution: {reduction.failure_code}" ) - if not expected_metrics.issubset(set(reduction.metrics)): + if set(reduction.metrics) != set(expected_reducer_metrics): raise AlgorithmExecutionError( - "Composite reducer did not return every declared metric" + "Composite reducer metrics do not match its declared metric surface" ) if "train_loss" not in reduction.metrics: raise AlgorithmExecutionError( @@ -1255,7 +923,7 @@ def _composite_backward( dist: Any, observation: dict[str, object] | None = None, metric_totals: dict[str, list[float]] | None = None, - expected_metrics: frozenset[str] = frozenset(), + expected_reducer_metrics: frozenset[str] = frozenset({"train_loss"}), ) -> object: """Reduce a Composite loss and return its Core-scaled backward scalar.""" reduction = _reduce_composite_loss( @@ -1265,7 +933,7 @@ def _composite_backward( device=device, dist=dist, observation=observation, - expected_metrics=expected_metrics, + expected_reducer_metrics=expected_reducer_metrics, ) if metric_totals is not None: _accumulate_metric_totals(metric_totals, reduction.metrics) @@ -1279,159 +947,6 @@ def _composite_backward( ) -def _restore_torch_retry_checkpoint( - checkpoint: object, - *, - stage_context: TorchStageContext, - model: object, - optimizer: object, - scheduler: object | None, - scaler: object, - rank: int, - strict_identity: bool = True, - progress_sink: dict[str, object] | None = None, - expected_accumulation: int | None = None, -) -> int: - """Validate a Ray-injected retry checkpoint before loading any state.""" - if checkpoint is None: - return 0 - identity = stage_context.runtime.run_identity - if identity is None: - raise AlgorithmExecutionError("Torch retry Stage context has no identity") - ref = TorchCheckpointRef(checkpoint) - descriptor = describe_torch_checkpoint( - ref, - object() - if not strict_identity - else TorchCheckpointContext( - stage=stage_context, - run_id=identity.run_id, - invocation_id=identity.invocation_id, - checkpoint_owner="core", - ), - ) - if strict_identity: - validate_torch_retry_identity( - descriptor, - identity, - world_size=stage_context.runtime.world_size, - ) - elif descriptor.world_size != stage_context.runtime.world_size: - raise AlgorithmExecutionError( - "Torch source checkpoint world size does not match current Stage" - ) - with _opened_checkpoint(checkpoint) as root: - import torch - - model_path = root / "model.pt" - optimizer_path = root / "optimizer.pt" - scaler_path = root / "scaler.pt" - rng_path = root / "rng_state.pt" - if any( - path.is_symlink() or not path.is_file() - for path in (model_path, optimizer_path, scaler_path, rng_path) - ): - raise AlgorithmExecutionError( - "Torch retry checkpoint is missing required model/optimizer/scaler/RNG state" - ) - target_model = cast(Any, getattr(model, "module", model)) - target_model.load_state_dict( - torch.load(model_path, map_location="cpu", weights_only=True) - ) - cast(Any, optimizer).load_state_dict( - torch.load(optimizer_path, map_location="cpu", weights_only=True) - ) - cast(Any, scaler).load_state_dict( - torch.load(scaler_path, map_location="cpu", weights_only=True) - ) - scheduler_path = root / "scheduler.pt" - if scheduler is not None: - if scheduler_path.is_symlink() or not scheduler_path.is_file(): - raise AlgorithmExecutionError( - "Torch retry checkpoint is missing configured scheduler state" - ) - cast(Any, scheduler).load_state_dict( - torch.load(scheduler_path, map_location="cpu", weights_only=True) - ) - payload = torch.load(rng_path, map_location="cpu", weights_only=True) - if not isinstance(payload, Mapping) or not isinstance( - payload.get("states"), list - ): - raise AlgorithmExecutionError("Torch retry RNG state is malformed") - states = payload["states"] - if ( - payload.get("world_size") != stage_context.runtime.world_size - or len(states) != stage_context.runtime.world_size - or rank >= len(states) - or not isinstance(states[rank], bytes) - ): - raise AlgorithmExecutionError( - "Torch retry checkpoint is missing the rank RNG state" - ) - rng = torch.frombuffer(states[rank], dtype=torch.uint8).clone() - torch.set_rng_state(rng) - cuda_states = payload.get("cuda_states_by_rank") - if torch.cuda.is_available(): - if ( - not isinstance(cuda_states, list) - or len(cuda_states) != stage_context.runtime.world_size - or rank >= len(cuda_states) - or not isinstance(cuda_states[rank], list) - or any(not isinstance(state, bytes) for state in cuda_states[rank]) - ): - raise AlgorithmExecutionError("Torch retry CUDA RNG state is malformed") - torch.cuda.set_rng_state_all( - [ - torch.frombuffer(state, dtype=torch.uint8).clone() - for state in cuda_states[rank] - ] - ) - progress_path = root / "torch_progress.json" - if progress_path.is_symlink() or not progress_path.is_file(): - raise AlgorithmExecutionError( - "Torch checkpoint is missing deterministic progress state" - ) - try: - progress = TorchCheckpointProgress.from_dict( - json.loads(progress_path.read_text(encoding="utf-8")) - ) - except (OSError, TypeError, ValueError) as exc: - raise AlgorithmExecutionError( - "Torch checkpoint progress state is malformed" - ) from exc - if progress.optimizer_step != descriptor.completed_step: - raise AlgorithmExecutionError( - "Torch checkpoint progress and descriptor step differ" - ) - if ( - expected_accumulation is not None - and progress.accumulation_steps != expected_accumulation - ): - raise AlgorithmExecutionError( - "Torch checkpoint accumulation configuration differs" - ) - if expected_accumulation is not None and ( - len(progress.dataset_cursor_by_rank) != stage_context.runtime.world_size - or str(rank) not in progress.dataset_cursor_by_rank - ): - raise AlgorithmExecutionError( - "Torch checkpoint dataset cursor does not cover every rank" - ) - if expected_accumulation is not None: - raw_rank_statistics = progress.rank_statistics - if ( - len(raw_rank_statistics) != stage_context.runtime.world_size - or str(rank) not in raw_rank_statistics - ): - raise AlgorithmExecutionError( - "Torch checkpoint statistics do not cover every rank" - ) - if progress_sink is not None: - progress_sink.update(progress.to_dict()) - progress_sink["_typed_progress"] = progress - return descriptor.completed_step - - def _finalize_torch_window( *, scaler: Any, @@ -1454,296 +969,40 @@ def _finalize_torch_window( optimizer.zero_grad() -def _should_apply_epoch_scheduler( - *, - restore_same_stage: bool, - epoch: int, - restored_epoch: int, - restored_epoch_scheduler_applied: bool, -) -> bool: - """Return whether the epoch boundary still owes one scheduler step.""" - return not ( - restore_same_stage - and epoch == restored_epoch - and restored_epoch_scheduler_applied - ) - - def _select_worker_checkpoint( config: Mapping[str, Any], stage_context: TorchStageContext, ) -> TorchWorkerCheckpointContext: - """Select retry first, then Core control, never the reverse.""" - import ray.train - - retry = ray.train.get_checkpoint() - identity = stage_context.runtime.run_identity - if identity is None: - raise AlgorithmExecutionError("Torch Worker Stage context has no identity") - if retry is not None: - descriptor = describe_torch_checkpoint( - TorchCheckpointRef(retry), - TorchCheckpointContext( - stage=stage_context, - run_id=identity.run_id, - invocation_id=identity.invocation_id, - checkpoint_owner="core", - ), - ) - validate_torch_retry_identity( - descriptor, - identity, - world_size=stage_context.runtime.world_size, - ) + """Expose only an invocation-local Stage dependency to the algorithm.""" + initial = config.get("_core_initial_checkpoint") + if initial is not None: + source = config.get("_core_checkpoint_source") + if source != "stage_dependency": + raise AlgorithmExecutionError("Torch initial checkpoint source is invalid") + descriptor = _describe_torch_checkpoint(TorchCheckpointRef(initial), object()) + if ( + descriptor.policy_digest != stage_context.runtime.policy_digest + or descriptor.execution_plan_digest + != stage_context.runtime.execution_plan_digest + or descriptor.world_size != stage_context.runtime.world_size + or descriptor.state_layout != stage_context.runtime.state_layout + ): + raise AlgorithmExecutionError("Torch initial checkpoint is incompatible") + if source == "stage_dependency" and ( + descriptor.identity.stage_id != stage_context.predecessor_stage_id + ): + raise AlgorithmExecutionError("Torch predecessor checkpoint is invalid") return TorchWorkerCheckpointContext( stage=stage_context, - source="ray_failure_retry", - checkpoint=TorchCheckpointRef( - retry, - descriptor_digest=descriptor.digest, - source_stage_id=descriptor.identity.stage_id, - descriptor=descriptor, - ), - ) - control_value = config.get("core_control") - if control_value is None: - return TorchWorkerCheckpointContext(stage=stage_context, source="none") - if not isinstance(control_value, Mapping): - raise AlgorithmExecutionError("Torch Core control envelope is malformed") - control = TorchWorkerControlEnvelope.from_dict(control_value) - if ( - control.target_stage_id != stage_context.stage_id - or control.run_id != identity.run_id - or control.invocation_id != identity.invocation_id - or control.policy_digest != stage_context.runtime.policy_digest - or control.execution_plan_digest != stage_context.runtime.execution_plan_digest - ): - raise AlgorithmExecutionError("Torch Core control envelope identity mismatch") - opener = config.get("_core_checkpoint_opener") - if opener is None: - opener_ref = config.get("_core_checkpoint_opener_ref") - if isinstance(opener_ref, str): - opener = _load_reference(QualifiedReference.parse(opener_ref)) - if not callable(opener): - raise AlgorithmExecutionError( - "Torch Core control envelope has no verified checkpoint opener" - ) - checkpoint = opener(control.checkpoint_locator) - descriptor = describe_torch_checkpoint( - TorchCheckpointRef(checkpoint), - object(), - ) - _require_checkpoint_commit(checkpoint, descriptor) - if descriptor.digest != control.checkpoint_descriptor_digest: - raise AlgorithmExecutionError("Torch Core control descriptor mismatch") - expected_binding = stage_context.runtime.input_binding_digest - if ( - descriptor.policy_digest != stage_context.runtime.policy_digest - or descriptor.execution_plan_digest - != stage_context.runtime.execution_plan_digest - or ( - expected_binding is not None - and descriptor.input_binding_digest != expected_binding - ) - or descriptor.identity.implementation_id - != stage_context.runtime.implementation_id - or descriptor.implementation_code_digest != identity.implementation_code_digest - or descriptor.world_size != stage_context.runtime.world_size - or descriptor.state_layout != stage_context.runtime.state_layout - ): - raise AlgorithmExecutionError("Torch Core control descriptor identity mismatch") - if control.purpose == "stage_dependency" and ( - descriptor.identity.stage_id != control.source_stage_id - or control.source_stage_id != stage_context.predecessor_stage_id - or descriptor.identity.run_id != identity.run_id - or descriptor.identity.invocation_id != identity.invocation_id - ): - raise AlgorithmExecutionError("Torch Core control source Stage mismatch") - if ( - control.purpose == "cross_run_initial_recovery" - and not descriptor.resume_supported - ): - raise AlgorithmExecutionError("Torch checkpoint is not externally recoverable") - return TorchWorkerCheckpointContext( - stage=stage_context, - source=( - "stage_dependency" - if control.purpose == "stage_dependency" - else "cross_run_initial_recovery" - ), - checkpoint=TorchCheckpointRef( - checkpoint, - descriptor_digest=descriptor.digest, - source_stage_id=descriptor.identity.stage_id, - descriptor=descriptor, - ), - ) - - -def open_torch_checkpoint_locator(locator: TorchCheckpointLocator) -> object: - """Core-owned locator opener hook; credentials are resolved by the runtime.""" - if not isinstance(locator, TorchCheckpointLocator): - raise AlgorithmExecutionError("Torch checkpoint locator is invalid") - if locator.uri.startswith("ray://"): - from ray.train import Checkpoint - - path = Path(locator.uri.removeprefix("ray://")) - if path.is_dir(): - try: - return Checkpoint.from_directory(str(path)) - except (OSError, ValueError, TypeError) as exc: - raise AlgorithmExecutionError( - "Torch checkpoint locator could not be opened" - ) from exc - if locator.uri.startswith(("s3://", "gs://", "gcs://", "hdfs://")): - from ray.train import Checkpoint - - try: - return Checkpoint(locator.uri) - except (OSError, ValueError, TypeError) as exc: - raise AlgorithmExecutionError( - "Torch checkpoint locator could not be opened" - ) from exc - raise AlgorithmExecutionError( - "Torch checkpoint locator requires a configured Core storage opener" - ) - - -def _persist_stage_checkpoint( - checkpoint: object, - *, - identity: TorchStageRunIdentity, - storage_path: object, - descriptor_digest: str, -) -> str | None: - """Persist a same-invocation Stage checkpoint without replacing prior data.""" - if not isinstance(storage_path, (str, Path)): - return None - raw_storage = str(storage_path) - if "://" in raw_storage and not raw_storage.startswith("file://"): - from urllib.parse import urlsplit - - opener = getattr(checkpoint, "as_directory", None) - if not callable(opener): - raise AlgorithmExecutionError( - "Torch remote Stage checkpoint cannot be opened for Core persistence" - ) - try: - import pyarrow.fs as pafs - - parsed = urlsplit(raw_storage) - filesystem, prefix = pafs.FileSystem.from_uri(raw_storage) - remote_staging = ( - f"{prefix.rstrip('/')}/{identity.run_config_name}/" - f".stage_checkpoint.staging-{descriptor_digest}" - ) - commit_path = f"{remote_staging}/torch_stage_commit.json" - commit_info = filesystem.get_file_info(commit_path) - commit_payload = { - "schema_version": 1, - "identity": identity.to_dict(), - "descriptor_digest": descriptor_digest, - } - if commit_info.type is pafs.FileType.File: - try: - existing_commit = json.loads( - filesystem.open_input_file(commit_path).read().decode("utf-8") - ) - except (OSError, TypeError, ValueError) as exc: - raise AlgorithmExecutionError( - "Torch remote Stage checkpoint commit marker is malformed" - ) from exc - if existing_commit != commit_payload: - raise AlgorithmExecutionError( - "Torch remote Stage checkpoint identity collision" - ) - return f"{parsed.scheme}://{parsed.netloc}/{remote_staging.lstrip('/')}" - filesystem.create_dir(remote_staging, recursive=True) - with opener() as source: - source_path = Path(source) - for source_file in source_path.rglob("*"): - if source_file.is_symlink(): - raise AlgorithmExecutionError( - "Torch Stage checkpoint payload contains a symlink" - ) - if not source_file.is_file(): - continue - relative = source_file.relative_to(source_path).as_posix() - destination = f"{remote_staging.rstrip('/')}/{relative}" - with filesystem.open_output_stream(destination) as stream: - stream.write(source_file.read_bytes()) - with filesystem.open_output_stream(commit_path) as stream: - stream.write( - (json.dumps(commit_payload, sort_keys=True) + "\n").encode() - ) - if not parsed.scheme or not parsed.netloc: - raise AlgorithmExecutionError("Torch remote storage URI is invalid") - return f"{parsed.scheme}://{parsed.netloc}/{remote_staging.lstrip('/')}" - except AlgorithmExecutionError: - raise - except (ImportError, OSError, TypeError, ValueError) as exc: - raise AlgorithmExecutionError( - "failed to persist Torch Stage checkpoint in Core storage" - ) from exc - root = Path(storage_path) - if not root.is_absolute(): - return None - target = root / identity.run_config_name / "stage_checkpoint" - commit_path_local = target / "torch_stage_commit.json" - commit_payload = { - "schema_version": 1, - "identity": identity.to_dict(), - "descriptor_digest": descriptor_digest, - } - if target.exists(): - if target.is_symlink() or not target.is_dir(): - raise AlgorithmExecutionError( - "Torch Stage checkpoint destination is not a directory" - ) - if not commit_path_local.is_file() or commit_path_local.is_symlink(): - raise AlgorithmExecutionError( - "Torch Stage checkpoint destination is a partial or uncommitted snapshot" - ) - try: - if ( - json.loads(commit_path_local.read_text(encoding="utf-8")) - != commit_payload - ): - raise AlgorithmExecutionError( - "Torch Stage checkpoint identity collision" - ) - except (OSError, TypeError, ValueError) as exc: - raise AlgorithmExecutionError( - "Torch Stage checkpoint commit marker is malformed" - ) from exc - return f"ray://{target}" - opener = getattr(checkpoint, "as_directory", None) - if not callable(opener): - return None - temporary_target = Path( - tempfile.mkdtemp( - prefix=f".{target.name}.staging-{descriptor_digest}-", - dir=target.parent, - ) - ) - with opener() as source: - source_path = Path(source) - for source_file in source_path.rglob("*"): - if source_file.is_symlink(): - raise AlgorithmExecutionError( - "Torch Stage checkpoint payload contains a symlink" - ) - if not source_file.is_file(): - continue - relative_local = source_file.relative_to(source_path) - destination_local = temporary_target / relative_local - destination_local.parent.mkdir(parents=True, exist_ok=True) - destination_local.write_bytes(source_file.read_bytes()) - (temporary_target / "torch_stage_commit.json").write_text( - json.dumps(commit_payload, sort_keys=True) + "\n", encoding="utf-8" + source=source, + checkpoint=TorchCheckpointRef( + initial, + descriptor_digest=descriptor.digest, + source_stage_id=descriptor.identity.stage_id, + descriptor=descriptor, + ), ) - os.replace(temporary_target, target) - return f"ray://{target}" + return TorchWorkerCheckpointContext(stage=stage_context, source="none") def _recipe_worker(config: Mapping[str, Any]) -> None: @@ -1757,7 +1016,14 @@ def _recipe_worker(config: Mapping[str, Any]) -> None: raise AlgorithmConfigurationError( "Core Worker implementation reference is missing" ) - recipe = _load_reference(QualifiedReference.parse(recipe_ref)) + recipe_reference = QualifiedReference.parse(recipe_ref) + recipe_digest = config.get("_core_implementation_code_digest") + if not isinstance(recipe_digest, str): + raise AlgorithmConfigurationError( + "Core Worker implementation code digest is missing" + ) + _validate_module_digest(recipe_reference, recipe_digest) + recipe = _load_reference(recipe_reference) if not isinstance(recipe, type) or not issubclass(recipe, TorchRecipe): raise AlgorithmConfigurationError("Core Worker did not receive a TorchRecipe") recipe_instance = recipe() @@ -1766,6 +1032,13 @@ def _recipe_worker(config: Mapping[str, Any]) -> None: raise AlgorithmConfigurationError("Core Worker stage context is missing") stage_context = TorchStageContext.from_dict(stage_context_value) runtime_context = stage_context.runtime + training = config.get("training", {}) + if not isinstance(training, Mapping): + raise AlgorithmConfigurationError("Torch training config must be a mapping") + seed = training.get("seed", 42) + if not isinstance(seed, int) or isinstance(seed, bool): + raise AlgorithmConfigurationError("Torch training seed must be an integer") + torch.manual_seed(seed) modules = recipe_instance.build_modules( TorchBuildContext(runtime=runtime_context, stage=stage_context) ) @@ -1812,9 +1085,6 @@ def _recipe_worker(config: Mapping[str, Any]) -> None: multi_role = len(training_roles) > 1 rank = ray.train.get_context().get_world_rank() world_size = ray.train.get_context().get_world_size() - training = config.get("training", {}) - if not isinstance(training, Mapping): - raise AlgorithmConfigurationError("Torch training config must be a mapping") epochs = training.get("epochs", 1) shuffle = training.get("shuffle", False) if not isinstance(shuffle, bool): @@ -1836,87 +1106,28 @@ def _recipe_worker(config: Mapping[str, Any]) -> None: if amp and not torch.cuda.is_available(): raise AlgorithmConfigurationError("Torch AMP requires CUDA") scaler = torch.amp.GradScaler("cuda", enabled=amp) - seed = training.get("seed", 42) - if not isinstance(seed, int) or isinstance(seed, bool): - raise AlgorithmConfigurationError("Torch training seed must be an integer") torch.manual_seed(seed + rank) - scheduler = optimization.scheduler accumulation = optimization.gradient_accumulation_steps checkpoint_context = _select_worker_checkpoint(config, stage_context) - restored_progress: dict[str, object] = {} - try: - loaded_step = _restore_torch_retry_checkpoint( - checkpoint_context.checkpoint.checkpoint - if checkpoint_context.checkpoint is not None - else None, - stage_context=stage_context, - model=model, - optimizer=optimization.optimizer, - scheduler=scheduler, - scaler=scaler, - rank=rank, - strict_identity=checkpoint_context.source == "ray_failure_retry", - progress_sink=restored_progress, - expected_accumulation=( - accumulation - if checkpoint_context.source - in {"ray_failure_retry", "cross_run_initial_recovery"} - else None - ), - ) - restored_step = ( - loaded_step - if checkpoint_context.source - in {"ray_failure_retry", "cross_run_initial_recovery"} - else 0 + if checkpoint_context.checkpoint is not None: + checkpoint_context.checkpoint.close() + if checkpoint_context.source != "none": + raise AlgorithmConfigurationError( + "TorchRecipe v1 accepts no external or predecessor checkpoint" ) - finally: - if checkpoint_context.checkpoint is not None: - checkpoint_context.checkpoint.close() - restore_same_stage = checkpoint_context.source in { - "ray_failure_retry", - "cross_run_initial_recovery", - } - typed_progress = restored_progress.get("_typed_progress") - if restore_same_stage and not isinstance(typed_progress, TorchCheckpointProgress): - raise AlgorithmExecutionError("Torch checkpoint progress is not typed") - restored_checkpoint_progress = ( - cast(TorchCheckpointProgress, typed_progress) if restore_same_stage else None - ) - rank_statistics = ( - restored_checkpoint_progress.rank_statistics.get(str(rank)) - if restored_checkpoint_progress is not None - else None - ) - if restore_same_stage and not isinstance( - rank_statistics, TorchRankProgressStatistics - ): - raise AlgorithmExecutionError("Torch checkpoint rank statistics are missing") - - rows = rank_statistics.rows_processed if rank_statistics is not None else 0 + # Ray owns Worker retry. Recipe retries deliberately restart the Stage from + # its deterministic seed and Dataset beginning; only the completed Stage is + # checkpointed below for export. + rows = 0 steps = 0 - coverage_totals = ( - dict(rank_statistics.coverage_totals) if rank_statistics is not None else {} - ) - loss_numerator_total = ( - rank_statistics.loss_numerator_total if rank_statistics is not None else 0.0 - ) - loss_normalizer_total = ( - rank_statistics.loss_normalizer_total if rank_statistics is not None else 0.0 - ) - metric_totals = ( - {name: list(pair) for name, pair in rank_statistics.metric_totals.items()} - if rank_statistics is not None - else {} - ) - evaluation_totals = ( - {name: list(pair) for name, pair in rank_statistics.evaluation_totals.items()} - if rank_statistics is not None - else {} - ) - reducer_observation: dict[str, object] = ( - dict(rank_statistics.reducer_observation) if rank_statistics is not None else {} - ) + optimizer_steps = 0 + coverage_totals: dict[str, int] = {} + loss_numerator_total = 0.0 + loss_normalizer_total = 0.0 + metric_totals: dict[str, list[float]] = {} + evaluation_totals: dict[str, list[float]] = {} + evaluation_rows: dict[str, int] = {} + reducer_observation: dict[str, object] = {} composite_loss_seen = False batch_context = TorchBatchContext( stage=stage_context, @@ -1937,241 +1148,8 @@ def _recipe_worker(config: Mapping[str, Any]) -> None: else None ), ) - checkpoint_interval = config.get("_core_checkpoint_interval_windows", 1) - if ( - not isinstance(checkpoint_interval, int) - or isinstance(checkpoint_interval, bool) - or checkpoint_interval < 1 - ): - raise AlgorithmConfigurationError( - "Torch checkpoint interval must be a positive integer" - ) import torch.distributed as dist - def emit_checkpoint( - completed_step: int, - *, - epoch: int, - micro_batch_cursor: int, - scheduler_step: int, - rows_processed: int, - coverage_totals: Mapping[str, int], - loss_numerator_total: float, - loss_normalizer_total: float, - metric_totals: Mapping[str, list[float]], - evaluation_totals: Mapping[str, list[float]], - reducer_observation: Mapping[str, object], - ) -> None: - """Report an optimizer-boundary checkpoint for Ray failure retry.""" - from ray.train import Checkpoint - - identity = runtime_context.run_identity - if identity is None: - raise AlgorithmExecutionError("Torch Worker stage context has no identity") - checkpoint_dir = Path(tempfile.mkdtemp(prefix="tributo_torch_checkpoint_")) - try: - target_model = getattr(model, "module", model) - torch.save(target_model.state_dict(), checkpoint_dir / "model.pt") - torch.save( - cast(Any, optimizer).state_dict(), checkpoint_dir / "optimizer.pt" - ) - torch.save(cast(Any, scaler).state_dict(), checkpoint_dir / "scaler.pt") - if scheduler is not None: - torch.save( - cast(Any, scheduler).state_dict(), checkpoint_dir / "scheduler.pt" - ) - cursor_by_rank: list[object] = [None] * world_size - if dist.is_available() and dist.is_initialized(): - dist.all_gather_object(cursor_by_rank, micro_batch_cursor) - else: - cursor_by_rank = [micro_batch_cursor] - if any( - not isinstance(value, int) or isinstance(value, bool) or value < 0 - for value in cursor_by_rank - ): - raise AlgorithmExecutionError( - "Torch dataset cursor collective is incomplete" - ) - local_statistics = TorchRankProgressStatistics( - rows_processed=rows_processed, - coverage_totals=coverage_totals, - loss_numerator_total=loss_numerator_total, - loss_normalizer_total=loss_normalizer_total, - metric_totals={ - name: (values[0], values[1]) - for name, values in metric_totals.items() - }, - evaluation_totals={ - name: (values[0], values[1]) - for name, values in evaluation_totals.items() - }, - reducer_observation=reducer_observation, - ) - statistics_by_rank: list[object] = [None] * world_size - if dist.is_available() and dist.is_initialized(): - dist.all_gather_object(statistics_by_rank, local_statistics.to_dict()) - else: - statistics_by_rank = [local_statistics.to_dict()] - if any(not isinstance(value, Mapping) for value in statistics_by_rank): - raise AlgorithmExecutionError( - "Torch checkpoint statistics collective is incomplete" - ) - progress = TorchCheckpointProgress( - epoch=epoch, - micro_batch_cursor=micro_batch_cursor, - optimizer_step=completed_step, - scheduler_step=scheduler_step, - accumulation_steps=accumulation, - dataset_cursor_by_rank={ - str(rank_id): cast(int, cursor) - for rank_id, cursor in enumerate(cursor_by_rank) - }, - shuffle_seed=int(seed + epoch), - rows_processed=rows_processed, - coverage_totals=coverage_totals, - loss_numerator_total=loss_numerator_total, - loss_normalizer_total=loss_normalizer_total, - metric_totals={ - name: (values[0], values[1]) - for name, values in metric_totals.items() - }, - evaluation_totals={ - name: (values[0], values[1]) - for name, values in evaluation_totals.items() - }, - rank_statistics={ - str(rank_id): TorchRankProgressStatistics.from_dict( - cast(Mapping[str, Any], stats) - ) - for rank_id, stats in enumerate(statistics_by_rank) - }, - epoch_scheduler_applied=False, - ) - (checkpoint_dir / "torch_progress.json").write_text( - json.dumps(progress.to_dict(), sort_keys=True, separators=(",", ":")), - encoding="utf-8", - ) - cpu_rng = torch.get_rng_state().cpu().numpy().tobytes() - cpu_states: list[bytes | None] = [None] * world_size - if dist.is_available() and dist.is_initialized(): - dist.all_gather_object(cpu_states, cpu_rng) - else: - cpu_states = [cpu_rng] - if any(not isinstance(value, bytes) for value in cpu_states): - raise AlgorithmExecutionError( - "Torch RNG state collective is incomplete" - ) - cuda_states_by_rank: list[list[bytes]] = [] - if torch.cuda.is_available(): - local_cuda = [ - state.cpu().numpy().tobytes() - for state in torch.cuda.get_rng_state_all() - ] - gathered_cuda: list[object] = [None] * world_size - if dist.is_available() and dist.is_initialized(): - dist.all_gather_object(gathered_cuda, local_cuda) - if any( - not isinstance(value, list) - or any(not isinstance(state, bytes) for state in value) - for value in gathered_cuda - ): - raise AlgorithmExecutionError( - "Torch CUDA RNG state collective is incomplete" - ) - cuda_states_by_rank = cast(list[list[bytes]], gathered_cuda) - else: - cuda_states_by_rank = [local_cuda] - torch.save( - { - "world_size": world_size, - "states": cast(list[bytes], cpu_states), - "cuda_states_by_rank": cuda_states_by_rank, - }, - checkpoint_dir / "rng_state.pt", - ) - payload_names = [ - "model.pt", - "optimizer.pt", - "scaler.pt", - "rng_state.pt", - "torch_progress.json", - ] - if scheduler is not None: - payload_names.append("scheduler.pt") - descriptor = TorchCheckpointDescriptor( - schema_version=1, - identity=identity, - run_config_name=identity.run_config_name, - state_layout=str(config.get("_core_state_layout", "replicated")), - world_size=world_size, - completed_step=completed_step, - policy_digest=runtime_context.policy_digest, - execution_plan_digest=runtime_context.execution_plan_digest, - input_binding_digest=str(config.get("_core_input_binding_digest", "")), - implementation_code_digest=str( - config.get("_core_implementation_code_digest", "") - ), - payload_files={ - name: hashlib.sha256( - (checkpoint_dir / name).read_bytes() - ).hexdigest() - for name in payload_names - }, - adapter_identity=config.get("_core_adapter_identity"), - resume_supported=runtime_context.resume_supported, - same_world_size_resume=runtime_context.same_world_size_resume, - ) - - class _IntervalDraft: - checkpoint_dir: str | os.PathLike[str] - - def __init__(self) -> None: - self.checkpoint_dir = str(checkpoint_dir) - - def report( - self, - *, - metrics: Mapping[str, object], - stage_context: object, - completed_step: int, - ) -> None: - del stage_context, completed_step - checkpoint = ( - Checkpoint.from_directory(str(checkpoint_dir)) - if rank == int(config.get("_core_checkpoint_owner_rank", 0)) - else None - ) - ray.train.report(dict(metrics), checkpoint=checkpoint) - - report_torch_checkpoint( - { - "train_loss": ( - loss_numerator_total / loss_normalizer_total - if loss_normalizer_total > 0 - else 0.0 - ), - "model_state_digest": hashlib.sha256( - json.dumps( - { - name: str( - cast(Any, value).detach().cpu().numpy().tobytes() - ) - for name, value in target_model.state_dict().items() - }, - sort_keys=True, - ).encode() - ).hexdigest(), - "checkpoint_descriptor": descriptor.to_dict(), - }, - _IntervalDraft(), - stage_context, - completed_step, - ) - finally: - import shutil - - shutil.rmtree(checkpoint_dir, ignore_errors=True) - def reduce_window(value: float) -> float: tensor = torch.tensor( value, dtype=torch.float64, device=next(model.parameters()).device @@ -2203,104 +1181,35 @@ def zeros(value: object) -> object: def aligned_metric_contributions( contributions: Mapping[str, TorchMetricContribution], ) -> dict[str, TorchMetricContribution]: - """Make metric keys identical before any later rank-wise reduction.""" + """Reject metric-key drift on every rank before later collectives.""" for _name, contribution in contributions.items(): if not isinstance(contribution, TorchMetricContribution): raise AlgorithmConfigurationError( "TorchStepResult metrics must be TorchMetricContribution values" ) - names = set(contributions) + observed = sorted(contributions) + expected = sorted(declared_metric_names) if dist.is_available() and dist.is_initialized(): gathered: list[object] = [None] * world_size - dist.all_gather_object(gathered, sorted(names)) - for value in gathered: - if not isinstance(value, list) or any( - not isinstance(name, str) for name in value - ): - raise AlgorithmExecutionError( - "Torch metric key collective is incomplete" - ) - names.update(value) - return { - name: contributions.get(name, TorchMetricContribution(0.0, 0.0)) - for name in sorted(names) - } + dist.all_gather_object(gathered, observed) + if any(value != expected for value in gathered): + raise AlgorithmExecutionError( + "TorchStepResult.metrics do not match TorchMetricPlan" + ) + elif observed != expected: + raise AlgorithmExecutionError( + "TorchStepResult.metrics do not match TorchMetricPlan" + ) + return dict(contributions) - raw_metric_mapping = config.get("_core_metric_mapping", {}) - metric_mapping = ( - {str(key): str(value) for key, value in raw_metric_mapping.items()} - if isinstance(raw_metric_mapping, Mapping) - else {} - ) - if len(set(metric_mapping.values())) != len(metric_mapping): - raise AlgorithmConfigurationError("Torch metric mapping targets must be unique") raw_metric_reducers = config.get("_core_metric_reducers", {}) if not isinstance(raw_metric_reducers, Mapping): raise AlgorithmConfigurationError("Torch metric reducers must be a mapping") declared_metric_names = {str(name) for name in raw_metric_reducers} - {"train_loss"} - expected_metric_names = frozenset( - { - source - for source, target in metric_mapping.items() - if target in declared_metric_names - } - | { - name - for name in declared_metric_names - if name not in metric_mapping.values() - } - ) - restored_epoch = ( - restored_checkpoint_progress.epoch - if restored_checkpoint_progress is not None - else 0 - ) - remaining_skip_micro_batches = ( - restored_checkpoint_progress.dataset_cursor_by_rank[str(rank)] - if restored_checkpoint_progress is not None - else 0 - ) - scheduler_steps = ( - restored_checkpoint_progress.scheduler_step - if restored_checkpoint_progress is not None - else 0 - ) - restored_epoch_scheduler_applied = ( - restored_checkpoint_progress.epoch_scheduler_applied - if restored_checkpoint_progress is not None - else False - ) - if restore_same_stage and restored_epoch >= epochs: - raise AlgorithmExecutionError( - "Torch checkpoint progress points past the configured epoch range" - ) - if ( - checkpoint_context.source - in { - "ray_failure_retry", - "cross_run_initial_recovery", - } - and restored_progress.get("shuffle_seed") != seed + restored_epoch - ): - raise AlgorithmExecutionError( - "Torch checkpoint shuffle state does not match the current seed" - ) - - def map_metric_contributions( - contributions: Mapping[str, TorchMetricContribution], - ) -> dict[str, TorchMetricContribution]: - return { - metric_mapping.get(name, name): contribution - for name, contribution in contributions.items() - } epoch_micro_batch_cursor = 0 for _epoch in range(epochs): - if _epoch < restored_epoch: - continue - epoch_micro_batch_cursor = ( - remaining_skip_micro_batches if _epoch == restored_epoch else 0 - ) + epoch_micro_batch_cursor = 0 next_payload: Any if multi_role: role_iterators = { @@ -2350,7 +1259,7 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: dist.all_reduce(first_active, op=dist.ReduceOp.SUM) if int(first_active.item()) == 0: continue - if remaining_skip_micro_batches == 0 and int(first_active.item()) != world_size: + if int(first_active.item()) != world_size: raise AlgorithmExecutionError( "TorchRecipe requires at least one batch on every rank" ) @@ -2360,7 +1269,7 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: "TorchRecipe.adapt_batch must return TorchBatch" ) window = TorchAccumulationWindow( - index=restored_step + steps // accumulation, + index=optimizer_steps, expected_micro_batches=accumulation, ) while True: @@ -2375,11 +1284,6 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: if int(active_tensor.item()) == 0: break next_raw = next_payload() - if remaining_skip_micro_batches > 0: - remaining_skip_micro_batches -= 1 - epoch_micro_batch_cursor += 1 - raw = next_raw - continue next_active = torch.tensor( 1 if next_raw is not None else 0, dtype=torch.int64, @@ -2428,8 +1332,8 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: batch, TorchStepContext( stage=stage_context, - window_index=restored_step + steps // accumulation, - micro_batch_index=steps % accumulation, + window_index=optimizer_steps, + micro_batch_index=window.observed_micro_batches, ), ) if not isinstance(step, TorchStepResult): @@ -2464,19 +1368,11 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: "TorchRecipe loss contribution is invalid" ) for metric_name, contribution in aligned_metric_contributions( - map_metric_contributions(step.metrics) + step.metrics ).items(): totals = metric_totals.setdefault(metric_name, [0.0, 0.0]) totals[0] += contribution.numerator totals[1] += contribution.normalizer - mapped_step_metrics = map_metric_contributions(step.metrics) - if ( - not isinstance(step.loss, TorchCompositeLossContribution) - and set(mapped_step_metrics) != declared_metric_names - ): - raise AlgorithmExecutionError( - "TorchStepResult.metrics do not match TorchMetricPlan" - ) backward_context = TorchBackwardContext( world_size=world_size, backward=lambda value: scaler.scale(value).backward(), @@ -2491,7 +1387,6 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: dist=dist, observation=reducer_observation, metric_totals=metric_totals, - expected_metrics=expected_metric_names, ) ), finalize_window=lambda scale: _finalize_torch_window( @@ -2511,6 +1406,7 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: backward_context, ) if result.window_complete: + optimizer_steps += 1 window = TorchAccumulationWindow( index=window.index + 1, expected_micro_batches=accumulation, @@ -2520,37 +1416,14 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: rows += batch.local_rows steps += 1 epoch_micro_batch_cursor += 1 - if ( - result.window_complete - and (steps // accumulation) % checkpoint_interval == 0 - ): - emit_checkpoint( - restored_step + steps // accumulation, - epoch=_epoch, - micro_batch_cursor=epoch_micro_batch_cursor, - scheduler_step=scheduler_steps, - rows_processed=rows, - coverage_totals=coverage_totals, - loss_numerator_total=loss_numerator_total, - loss_normalizer_total=loss_normalizer_total, - metric_totals=metric_totals, - evaluation_totals=evaluation_totals, - reducer_observation=reducer_observation, - ) raw = next_raw - if optimization.scheduler is not None and _should_apply_epoch_scheduler( - restore_same_stage=restore_same_stage, - epoch=_epoch, - restored_epoch=restored_epoch, - restored_epoch_scheduler_applied=restored_epoch_scheduler_applied, - ): + if optimization.scheduler is not None: scheduler_step = getattr(optimization.scheduler, "step", None) if not callable(scheduler_step): raise AlgorithmConfigurationError( "Torch scheduler must implement step()" ) scheduler_step() - scheduler_steps += 1 for split in ("val", "test"): try: evaluation_data = ray.train.get_dataset_shard(split) @@ -2621,16 +1494,9 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: raise AlgorithmConfigurationError( "TorchRecipe.validation_step returned invalid result" ) - mapped_validation_metrics = map_metric_contributions(validation.metrics) - if ( - not isinstance(validation.loss, TorchCompositeLossContribution) - and set(mapped_validation_metrics) != declared_metric_names - ): - raise AlgorithmExecutionError( - "TorchRecipe.validation_step.metrics do not match TorchMetricPlan" - ) + evaluation_rows[split] = evaluation_rows.get(split, 0) + batch.local_rows for metric_name, contribution in aligned_metric_contributions( - mapped_validation_metrics + validation.metrics ).items(): _accumulate_metric_totals( evaluation_totals, @@ -2649,7 +1515,6 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: device=next(model.parameters()).device, dist=dist, observation=reducer_observation, - expected_metrics=expected_metric_names, ) evaluation_metrics = { name: contribution @@ -2697,6 +1562,12 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: node_id = ray.get_runtime_context().get_node_id() input_rows_evidence = dict(coverage_totals) input_rows_evidence[stage_roles[0]] = rows + for split, split_rows in evaluation_rows.items(): + if split in input_rows_evidence: + raise AlgorithmExecutionError( + f"Torch evaluation role {split!r} conflicts with coverage evidence" + ) + input_rows_evidence[split] = split_rows execution_worker = { "worker_id": f"torch-{rank}", "node_id": str(node_id), @@ -2718,186 +1589,11 @@ def _next_single_payload(iterator: Any = iterator) -> object | None: if any(not isinstance(item, Mapping) for item in worker_records_raw): raise AlgorithmExecutionError("Torch worker evidence collective is incomplete") worker_records = [cast(dict[str, object], item) for item in worker_records_raw] - from ray.train import Checkpoint - checkpoint_dir = Path(tempfile.mkdtemp(prefix="tributo_torch_checkpoint_")) try: torch.save( getattr(model, "module", model).state_dict(), checkpoint_dir / "model.pt" ) - torch.save(cast(Any, optimizer).state_dict(), checkpoint_dir / "optimizer.pt") - torch.save(cast(Any, scaler).state_dict(), checkpoint_dir / "scaler.pt") - if optimization.scheduler is not None: - torch.save( - cast(Any, optimization.scheduler).state_dict(), - checkpoint_dir / "scheduler.pt", - ) - final_cursor_values: list[object] = [None] * world_size - if dist.is_available() and dist.is_initialized(): - dist.all_gather_object(final_cursor_values, epoch_micro_batch_cursor) - else: - final_cursor_values = [epoch_micro_batch_cursor] - if any( - not isinstance(value, int) or isinstance(value, bool) or value < 0 - for value in final_cursor_values - ): - raise AlgorithmExecutionError("Torch final cursor collective is incomplete") - final_local_statistics = TorchRankProgressStatistics( - rows_processed=rows, - coverage_totals=coverage_totals, - loss_numerator_total=loss_numerator_total, - loss_normalizer_total=loss_normalizer_total, - metric_totals={ - name: (values[0], values[1]) for name, values in metric_totals.items() - }, - evaluation_totals={ - name: (values[0], values[1]) - for name, values in evaluation_totals.items() - }, - reducer_observation=reducer_observation, - ) - final_statistics_by_rank: list[object] = [None] * world_size - if dist.is_available() and dist.is_initialized(): - dist.all_gather_object( - final_statistics_by_rank, final_local_statistics.to_dict() - ) - else: - final_statistics_by_rank = [final_local_statistics.to_dict()] - if any(not isinstance(value, Mapping) for value in final_statistics_by_rank): - raise AlgorithmExecutionError( - "Torch final statistics collective is incomplete" - ) - final_progress = TorchCheckpointProgress( - epoch=max(epochs - 1, 0), - micro_batch_cursor=epoch_micro_batch_cursor, - optimizer_step=restored_step + steps // accumulation, - scheduler_step=scheduler_steps, - accumulation_steps=accumulation, - dataset_cursor_by_rank={ - str(rank_id): cast(int, cursor) - for rank_id, cursor in enumerate(final_cursor_values) - }, - shuffle_seed=int(seed + max(epochs - 1, 0)), - rows_processed=rows, - coverage_totals=coverage_totals, - loss_numerator_total=loss_numerator_total, - loss_normalizer_total=loss_normalizer_total, - metric_totals={ - name: (values[0], values[1]) for name, values in metric_totals.items() - }, - evaluation_totals={ - name: (values[0], values[1]) - for name, values in evaluation_totals.items() - }, - rank_statistics={ - str(rank_id): TorchRankProgressStatistics.from_dict( - cast(Mapping[str, Any], stats) - ) - for rank_id, stats in enumerate(final_statistics_by_rank) - }, - epoch_scheduler_applied=True, - ) - (checkpoint_dir / "torch_progress.json").write_text( - json.dumps(final_progress.to_dict(), sort_keys=True, separators=(",", ":")), - encoding="utf-8", - ) - rng_payload = torch.get_rng_state().cpu().numpy().tobytes() - rng_states: list[bytes | None] = [None] * world_size - if dist.is_available() and dist.is_initialized(): - dist.all_gather_object(rng_states, rng_payload) - else: - rng_states = [rng_payload] - if any(not isinstance(value, bytes) for value in rng_states): - raise AlgorithmExecutionError("Torch RNG state collective is incomplete") - cuda_rng_payload: list[list[bytes]] = [] - if torch.cuda.is_available(): - local_cuda = [ - state.cpu().numpy().tobytes() - for state in torch.cuda.get_rng_state_all() - ] - if dist.is_available() and dist.is_initialized(): - all_cuda: list[object] = [None] * world_size - dist.all_gather_object(all_cuda, local_cuda) - if any( - not isinstance(value, list) - or any(not isinstance(state, bytes) for state in value) - for value in all_cuda - ): - raise AlgorithmExecutionError( - "Torch CUDA RNG state collective is incomplete" - ) - cuda_rng_payload = cast(list[list[bytes]], all_cuda) - else: - cuda_rng_payload = [local_cuda] - torch.save( - { - "world_size": world_size, - "states": rng_states, - "cuda_states_by_rank": cuda_rng_payload, - }, - checkpoint_dir / "rng_state.pt", - ) - identity = stage_context.runtime.run_identity - if identity is None: - raise AlgorithmExecutionError( - "Torch Worker stage context has no run identity" - ) - payload_names = [ - "model.pt", - "optimizer.pt", - "scaler.pt", - "rng_state.pt", - "torch_progress.json", - ] - if optimization.scheduler is not None: - payload_names.append("scheduler.pt") - payload_files = { - name: hashlib.sha256((checkpoint_dir / name).read_bytes()).hexdigest() - for name in payload_names - } - descriptor = TorchCheckpointDescriptor( - schema_version=1, - identity=identity, - run_config_name=torch_run_config_name(identity), - state_layout=str(config.get("_core_state_layout", "replicated")), - world_size=world_size, - completed_step=restored_step + steps // accumulation, - policy_digest=stage_context.runtime.policy_digest, - execution_plan_digest=stage_context.runtime.execution_plan_digest, - input_binding_digest=str(config.get("_core_input_binding_digest", "")), - implementation_code_digest=str( - config.get("_core_implementation_code_digest", "") - ), - payload_files=payload_files, - adapter_identity=config.get("_core_adapter_identity"), - resume_supported=stage_context.runtime.resume_supported, - same_world_size_resume=stage_context.runtime.same_world_size_resume, - ) - - class _Draft: - checkpoint_dir: str | os.PathLike[str] - - def __init__( - self, checkpoint_dir: Path, checkpoint_owner_rank: int - ) -> None: - self.checkpoint_dir = str(checkpoint_dir) - self._checkpoint_owner_rank = checkpoint_owner_rank - - def report( - self, - *, - metrics: Mapping[str, object], - stage_context: object, - completed_step: int, - ) -> None: - del stage_context, completed_step - checkpoint = ( - Checkpoint.from_directory(str(checkpoint_dir)) - if rank == self._checkpoint_owner_rank - else None - ) - ray.train.report(dict(metrics), checkpoint=checkpoint) - loss_state = torch.tensor( [loss_numerator_total, loss_normalizer_total], dtype=torch.float64, @@ -2905,6 +1601,10 @@ def report( ) if dist.is_available() and dist.is_initialized(): dist.all_reduce(loss_state, op=dist.ReduceOp.SUM) + if not composite_loss_seen and loss_state[1].item() <= 0: + raise AlgorithmExecutionError( + "Torch training produced no positive loss normalizer" + ) train_loss = ( float(loss_state[0].item() / loss_state[1].item()) if loss_state[1].item() > 0 @@ -2916,10 +1616,19 @@ def report( ) if not isinstance(metric_reducers, Mapping): raise AlgorithmConfigurationError("Torch metric reducers must be a mapping") + resolved_metric_reducers = { + str(name): str(reducer) for name, reducer in metric_reducers.items() + } + for split in ("val", "test"): + for name, reducer in metric_reducers.items(): + resolved_metric_reducers[f"{split}_{name}"] = str(reducer) + resolved_metric_reducers[f"{split}_loss"] = str( + metric_reducers["train_loss"] + ) metric_values.update( _reduce_metric_totals( metric_totals, - metric_reducers, + resolved_metric_reducers, device=next(model.parameters()).device, dist=dist, world_size=world_size, @@ -2932,16 +1641,13 @@ def report( metric_values.update( _reduce_metric_totals( evaluation_totals, - metric_reducers, + resolved_metric_reducers, device=next(model.parameters()).device, dist=dist, world_size=world_size, ) ) metric_values.setdefault("train_loss", train_loss) - for source_name, target_name in metric_mapping.items(): - if source_name in metric_values: - metric_values[target_name] = metric_values.pop(source_name) reducer_report = { "reducer_id": config.get("_core_global_loss_reducer_id"), "reducer_api_version": config.get("_core_global_loss_reducer_api_version"), @@ -2959,15 +1665,14 @@ def report( **metric_values, "execution_workers": worker_records, "model_state_digest": state_digest, - "checkpoint_descriptor": descriptor.to_dict(), **reducer_report, }, - _Draft( - checkpoint_dir, - int(config.get("_core_checkpoint_owner_rank", 0)), + TorchCheckpointPayloadDraft( + checkpoint_dir=checkpoint_dir, + checkpoint_owner_rank=int(config.get("_core_checkpoint_owner_rank", 0)), ), stage_context, - restored_step + steps // accumulation, + optimizer_steps, ) finally: import shutil @@ -2977,19 +1682,26 @@ def report( @DeveloperAPI def torch_recipe_train_loop_per_worker(config: Mapping[str, Any]) -> None: - """Public Core worker entrypoint referenced by ``TorchStageSpec``.""" + """Run the Core-selected Recipe Worker wrapper.""" _recipe_worker(config) @DeveloperAPI def ray_torch_adapter_train_loop_per_worker(config: Mapping[str, Any]) -> None: - """Public Core wrapper entrypoint for Adapter-owned Stage loops.""" + """Run the Core-selected Adapter Worker wrapper.""" reference = config.get("_core_implementation_ref") if not isinstance(reference, str): raise AlgorithmConfigurationError( "Adapter Worker implementation reference is missing" ) - implementation = _load_reference(QualifiedReference.parse(reference)) + implementation_reference = QualifiedReference.parse(reference) + implementation_digest = config.get("_core_implementation_code_digest") + if not isinstance(implementation_digest, str): + raise AlgorithmConfigurationError( + "Adapter Worker implementation code digest is missing" + ) + _validate_module_digest(implementation_reference, implementation_digest) + implementation = _load_reference(implementation_reference) if not isinstance(implementation, type): raise AlgorithmConfigurationError("Adapter Worker reference is invalid") adapter = implementation() @@ -3022,58 +1734,17 @@ def preflight( self, plan: Any, run_id: str, - invocation_id: str, - ) -> TorchPreflightLease: + ) -> None: policy = _policy(plan) - if policy.state_layout == "sharded": - raise AlgorithmConfigurationError( - "Torch sharded state is reserved and not supported by Runtime v1" - ) - if policy.evidence_adapter_ref is not None: - raise AlgorithmConfigurationError( - "Torch evidence_adapter_ref is reserved until a Core evidence adapter protocol is gated" - ) - ray_config = plan.algorithm_config.get("ray", {}) - resume_config = ( - ray_config.get("resume", {}) if isinstance(ray_config, Mapping) else {} - ) - has_external_recovery = plan.runtime.resume_from is not None or ( - isinstance(resume_config, Mapping) - and any( - resume_config.get(name) is not None - for name in ("uri", "checkpoint_uri", "checkpoint_descriptor_digest") - ) - ) - if plan.runtime.torch_recovery is not None: - recovery = TorchRecoveryEnvelope.from_dict(plan.runtime.torch_recovery) - has_external_recovery = has_external_recovery or bool( - recovery.stage_checkpoints or recovery.active_checkpoint is not None - ) + _torch_ray_config(plan) + _validate_runtime_policy(policy) + has_external_recovery = plan.runtime.resume_from is not None if has_external_recovery and not policy.resume_supported: raise AlgorithmConfigurationError( "Torch Policy does not support external recovery" ) implementation = _load_torch_implementation(plan) - for stage in policy.execution_plan.stages: - expected_loop_ref = ( - _CORE_ADAPTER_LOOP_REF - if policy.loop_owner == "adapter" - else _CORE_RECIPE_LOOP_REF - ) - if stage.worker_loop_ref != expected_loop_ref: - raise AlgorithmConfigurationError( - "Torch Stage worker_loop_ref must use the Core-owned loop wrapper" - ) - stage_worker = _load_reference( - QualifiedReference.parse(stage.worker_loop_ref) - ) - if not callable(stage_worker): - raise AlgorithmConfigurationError( - f"Torch Stage {stage.stage_id!r} worker_loop_ref is not callable" - ) - identity = _identity( - plan, run_id, invocation_id, policy.execution_plan.final_stage_id - ) + identity = _identity(plan, run_id, run_id, policy.execution_plan.final_stage_id) context = TorchRuntimeContext( algorithm_config=_torch_algorithm_context_config(plan), implementation_id=plan.implementation.implementation_id, @@ -3091,7 +1762,6 @@ def preflight( else None ), resume_supported=policy.resume_supported, - same_world_size_resume=policy.same_world_size_resume, ) if isinstance(implementation, RayTorchAdapter): implementation.validate_environment(context) @@ -3144,19 +1814,6 @@ def preflight( raise AlgorithmConfigurationError( "Torch global reducer code digest mismatch" ) - token = TorchPreflightTokenData( - run_id=run_id, - invocation_id=invocation_id, - algorithm=plan.resolution.algorithm, - implementation_ref=str(plan.implementation.implementation_ref), - implementation_code_digest=cast(str, plan.implementation.code_digest), - policy_digest=policy.digest, - execution_plan_digest=policy.execution_plan.digest, - runtime_id=self.runtime_id, - reducer_identity=policy.global_loss_reducer_ref, - plan_digest=plan.plan_id, - ) - return TorchPreflightLease(token) @staticmethod def _validate_recipe_environment(context: TorchRuntimeContext) -> None: @@ -3173,70 +1830,35 @@ def _validate_recipe_environment(context: TorchRuntimeContext) -> None: if not hasattr(ray, "train") or not hasattr(torch, "nn"): raise AlgorithmConfigurationError("TorchRecipe environment is incomplete") - def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionResult: - if not isinstance(envelope, TorchRuntimeExecutionEnvelope): - raise AlgorithmConfigurationError( - "Ray Train Torch requires TorchRuntimeExecutionEnvelope" - ) - base = envelope.base + def execute(self, envelope: RuntimeExecutionEnvelope) -> WorkerExecutionResult: + base = envelope if base.cancelled: raise AlgorithmExecutionError("Torch execution was cancelled") - invocation_id = envelope.preflight_lease.data.invocation_id - token = envelope.preflight_lease.consume( - run_id=base.run_id, - invocation_id=invocation_id, - plan_digest=base.plan.plan_id, - runtime_id=self.runtime_id, - ) + invocation_id = base.run_id implementation = _load_torch_implementation(base.plan) policy = _policy(base.plan) - if policy.state_layout == "sharded": - raise AlgorithmConfigurationError( - "Torch sharded state is reserved and not supported by Runtime v1" - ) - if policy.evidence_adapter_ref is not None: + _validate_runtime_policy(policy) + reducer_metadata = _reducer_metadata(policy) + resume_from = base.plan.runtime.resume_from + if resume_from is not None: raise AlgorithmConfigurationError( - "Torch evidence_adapter_ref is reserved until a Core evidence adapter protocol is gated" + "Torch Runtime API v1 does not support cross-Run recovery" ) - reducer_metadata = _reducer_metadata(policy) - completed_stage_ids, active_stage_id, recovery_records = _recovery_records( - base.plan, - policy, - worker_count=base.plan.runtime.worker_count, - ) prepared = _prepare_datasets(base) + configured_storage_path, max_failures = _torch_ray_config(base.plan) + owned_storage_dirs: list[tempfile.TemporaryDirectory[str]] = [] try: stages = policy.execution_plan.stages - stage_records: dict[str, dict[str, Any]] = dict(recovery_records) + stage_records: dict[str, dict[str, Any]] = {} last_result: Any = None final_result: Any = None stage_evidence: list[ComponentStageEvidence] = [] + stage_metrics: dict[str, float] = {} final_expected_rows: Mapping[str, int] = {} + final_replicated_bytes: Mapping[str, int] = {} for index, stage in enumerate(stages): - if stage.stage_id in completed_stage_ids: - recovered_record = stage_records.get(stage.stage_id) - if recovered_record is None: - raise AlgorithmExecutionError( - f"Torch recovery is missing Stage {stage.stage_id!r}" - ) - recovered_evidence = _recovered_stage_evidence( - plan=base.plan, - policy=policy, - stage=stage, - descriptor=recovered_record["descriptor"], - evidence=recovered_record.get("evidence", {}), - ) - if policy.state_layout == "component": - stage_evidence.append(recovered_evidence) - if stage.stage_id == policy.execution_plan.final_stage_id: - final_expected_rows = { - role.role: int(role.expected_rows or 0) - for role in recovered_evidence.roles - if role.present and role.expected_rows is not None - } - continue identity = _identity( - base.plan, token.run_id, token.invocation_id, stage.stage_id + base.plan, base.run_id, invocation_id, stage.stage_id ) runtime_context = TorchRuntimeContext( algorithm_config=_torch_algorithm_context_config(base.plan), @@ -3255,7 +1877,6 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes else None ), resume_supported=policy.resume_supported, - same_world_size_resume=policy.same_world_size_resume, ) predecessor_id = stage.checkpoint_from_stage predecessor_record = ( @@ -3290,7 +1911,8 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes raise AlgorithmConfigurationError( "RayTorchAdapter.bind_datasets must return named datasets" ) - expected_rows = _validate_stage_routes( + stage_datasets = dict(stage_datasets) + expected_rows, replicated_bytes_by_role = _validate_stage_routes( policy, stage, stage_datasets, @@ -3298,6 +1920,7 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes ) if stage.stage_id == policy.execution_plan.final_stage_id: final_expected_rows = dict(expected_rows) + final_replicated_bytes = dict(replicated_bytes_by_role) if ( stage.checkpoint_from_stage is not None and predecessor_record is None @@ -3306,31 +1929,13 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes f"Torch Stage {stage.stage_id!r} is missing checkpoint from " f"{stage.checkpoint_from_stage!r}" ) - core_control = _control_for_stage( - base.plan, - policy, - stage, - run_id=token.run_id, - invocation_id=token.invocation_id, - checkpoint=( - stage_records.get(stage.stage_id) - if active_stage_id == stage.stage_id - else stage_records.get(stage.checkpoint_from_stage) - if stage.checkpoint_from_stage is not None - else None - ), - purpose=( - "cross_run_initial_recovery" - if active_stage_id == stage.stage_id - else "stage_dependency" - if stage.checkpoint_from_stage is not None - else None - ), - source_stage_id=( - None - if active_stage_id == stage.stage_id - else stage.checkpoint_from_stage - ), + initial_checkpoint = ( + stage_records[stage.checkpoint_from_stage]["checkpoint"] + if stage.checkpoint_from_stage is not None + else None + ) + checkpoint_source = ( + "stage_dependency" if initial_checkpoint is not None else "none" ) train_config = _torch_algorithm_context_config(base.plan) stage_binding = base.plan.input_bindings.get(stage.input_roles[0]) @@ -3352,9 +1957,6 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes "_core_weight_name": stage_binding.sample_weight_name, "_core_stage_input_roles": list(stage.input_roles), "_core_input_role_bindings": _torch_input_bindings(base.plan), - "_core_checkpoint_interval_windows": int( - getattr(stage, "checkpoint_interval_windows", 1) - ), "_core_policy_digest": policy.digest, "_core_execution_plan_digest": policy.execution_plan.digest, "_core_global_loss_reducer_ref": policy.global_loss_reducer_ref, @@ -3371,15 +1973,12 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes name: reduction.value for name, reduction in policy.metric_reducers.items() }, - "_core_metric_mapping": dict( - getattr(stage, "metric_mapping", {}) - ), "_core_adapter_identity": ( base.plan.implementation.implementation_id if isinstance(implementation, RayTorchAdapter) else None ), - "_core_stage_context": context.to_dict(), + "_core_stage_context": _worker_stage_context(context).to_dict(), "_core_batch_size": int( base.plan.algorithm_config.get("training", {}).get( "batch_size", 32 @@ -3390,25 +1989,13 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes else 32 ), "_core_torch_evidence": {}, - "core_control": core_control, - "_core_checkpoint_opener_ref": ( - "tributo.integrations.algorithm_runtimes.ray_train_torch:" - "open_torch_checkpoint_locator" - if core_control is not None - else None - ), + "_core_initial_checkpoint": initial_checkpoint, + "_core_checkpoint_source": checkpoint_source, } ) loop: Any if isinstance(implementation, TorchRecipe): - loop_ref = _load_reference( - QualifiedReference.parse(stage.worker_loop_ref) - ) - if not callable(loop_ref): - raise AlgorithmConfigurationError( - "Recipe Stage worker_loop_ref is not callable" - ) - loop = loop_ref + loop = torch_recipe_train_loop_per_worker else: adapter = cast(RayTorchAdapter, implementation) adapter_config = adapter.worker_config(context) @@ -3428,14 +2015,7 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes "Adapter worker config must be JSON-compatible" ) from exc train_config["adapter_config"] = dict(adapter_config) - loop_ref = _load_reference( - QualifiedReference.parse(stage.worker_loop_ref) - ) - if not callable(loop_ref): - raise AlgorithmConfigurationError( - "Adapter Stage worker_loop_ref is not callable" - ) - loop = loop_ref + loop = ray_torch_adapter_train_loop_per_worker from ray.train import FailureConfig, RunConfig, ScalingConfig from ray.train.torch import TorchConfig, TorchTrainer @@ -3443,28 +2023,16 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes TorchRoleDataConfig, ) - storage_path = None - ray_config = base.plan.algorithm_config.get("ray", {}) - max_failures = 0 - if isinstance(ray_config, Mapping): - storage_path = ray_config.get("storage_path") - configured_failures = ray_config.get("max_failures", 0) - if ( - not isinstance(configured_failures, int) - or isinstance(configured_failures, bool) - or configured_failures < -1 - ): - raise AlgorithmConfigurationError( - "ray.max_failures must be -1 or a non-negative integer" - ) - max_failures = configured_failures + storage_path = configured_storage_path if ( storage_path is None and cast(Any, base.plan.runtime.execution_profile).value == "local" ): - storage_path = tempfile.mkdtemp(prefix="tributo_torch_runs_") - if storage_path is not None: - claim_torch_run_directory(storage_path, identity) + owned_storage = tempfile.TemporaryDirectory( + prefix="tributo_torch_runs_" + ) + owned_storage_dirs.append(owned_storage) + storage_path = owned_storage.name trainer = TorchTrainer( train_loop_per_worker=loop, train_loop_config=train_config, @@ -3476,7 +2044,7 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes ), datasets=cast(Any, dict(stage_datasets)), run_config=RunConfig( - name=torch_run_config_name(identity), + name=identity.run_config_name, storage_path=str(storage_path) if storage_path is not None else None, @@ -3493,12 +2061,18 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes if stage.stage_id == policy.execution_plan.final_stage_id: final_result = last_result metrics = last_result.metrics or {} + _record_stage_metrics( + stage_metrics, + metrics, + set(policy.metric_reducers), + is_final=stage.stage_id == policy.execution_plan.final_stage_id, + ) stage_descriptor_payload = metrics.get("checkpoint_descriptor") checkpoint = getattr(last_result, "checkpoint", None) if checkpoint is not None and isinstance( stage_descriptor_payload, Mapping ): - validated_descriptor = describe_torch_checkpoint( + validated_descriptor = _describe_torch_checkpoint( TorchCheckpointRef(checkpoint), TorchCheckpointContext( stage=context, @@ -3511,24 +2085,8 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes raise AlgorithmExecutionError( "Torch Stage checkpoint descriptor differs from embedded payload" ) - persisted_locator = _persist_stage_checkpoint( - checkpoint, - identity=identity, - storage_path=storage_path, - descriptor_digest=validated_descriptor.digest, - ) - if persisted_locator is not None: - stage_descriptor_payload = dict(stage_descriptor_payload) - stage_descriptor_payload["locator"] = persisted_locator - stage_descriptor_payload["descriptor_digest"] = ( - validated_descriptor.digest - ) stage_records[stage.stage_id] = { - "locator": ( - stage_descriptor_payload.get("locator") - if isinstance(stage_descriptor_payload, Mapping) - else None - ), + "checkpoint": checkpoint, "descriptor_digest": validated_descriptor.digest, "descriptor": validated_descriptor.to_dict(), "evidence": { @@ -3546,7 +2104,7 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes if key in metrics }, } - elif stage.checkpoint_required: + else: raise AlgorithmExecutionError( f"Torch Stage {stage.stage_id!r} requires a Core checkpoint" ) @@ -3558,6 +2116,7 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes identity=identity, metrics=metrics, expected_rows=expected_rows, + replicated_bytes_by_role=replicated_bytes_by_role, ) ) final_stage = next( @@ -3565,27 +2124,18 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes for stage in stages if stage.stage_id == policy.execution_plan.final_stage_id ) - if final_result is None and final_stage.stage_id in stage_records: - recovered = stage_records[final_stage.stage_id] - locator = TorchCheckpointLocator( - cast(str, recovered["locator"]), - cast(str, recovered["descriptor_digest"]), - ) - recovered_checkpoint = open_torch_checkpoint_locator(locator) - final_result = _CheckpointResultProxy( - recovered_checkpoint, - recovered_checkpoint, - metrics={ - **dict(recovered.get("evidence", {})), - "checkpoint_descriptor": recovered["descriptor"], - }, - ) if final_result is None: raise AlgorithmExecutionError( "Torch execution plan has no Stage result" ) last_result = final_result metrics = dict(final_result.metrics or {}) + metrics.update(stage_metrics) + missing_metrics = sorted(set(policy.metric_reducers) - set(stage_metrics)) + if missing_metrics: + raise AlgorithmExecutionError( + f"Torch execution did not report declared metric(s): {missing_metrics}" + ) checkpoint = getattr(final_result, "checkpoint", None) if checkpoint is not None: raw_descriptor = metrics.get("checkpoint_descriptor") @@ -3600,7 +2150,7 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes source_stage_id=descriptor.identity.stage_id, descriptor=descriptor, ) - describe_torch_checkpoint( + _describe_torch_checkpoint( checkpoint_ref, TorchCheckpointContext( stage=_stage_context( @@ -3624,7 +2174,6 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes else None ), resume_supported=policy.resume_supported, - same_world_size_resume=policy.same_world_size_resume, ), final_stage, stages.index(final_stage), @@ -3636,8 +2185,8 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes ) final_identity = _identity( base.plan, - token.run_id, - token.invocation_id, + base.run_id, + invocation_id, policy.execution_plan.final_stage_id, ) supplied_evidence = metrics.get("torch_evidence") @@ -3663,6 +2212,7 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes stage=final_stage, workers=workers, expected_rows=final_expected_rows, + replicated_bytes_by_role=final_replicated_bytes, ) global_digest = metrics.get("model_state_digest") if not isinstance(global_digest, str) or len(global_digest) != 64: @@ -3680,7 +2230,7 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes ).hexdigest() metrics["torch_evidence"] = TorchExecutionEvidence( identity=final_identity, - run_config_name=torch_run_config_name(final_identity), + run_config_name=final_identity.run_config_name, policy_digest=policy.digest, parallelism_id=policy.parallelism_id, state_layout=policy.state_layout, @@ -3764,13 +2314,12 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes else None ), resume_supported=policy.resume_supported, - same_world_size_resume=policy.same_world_size_resume, ), final_stage, stages.index(final_stage), ), - run_id=token.run_id, - invocation_id=token.invocation_id, + run_id=base.run_id, + invocation_id=invocation_id, checkpoint_owner="core", ) checkpoint = implementation.checkpoint_source( @@ -3785,7 +2334,7 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes if isinstance(checkpoint, TorchCheckpointRef) else TorchCheckpointRef(checkpoint) ) - describe_torch_checkpoint(checkpoint_ref, final_context) + _describe_torch_checkpoint(checkpoint_ref, final_context) export_result = _CheckpointResultProxy( last_result, checkpoint_ref.checkpoint, @@ -3795,7 +2344,7 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes execution = export_ray_train_torch_result( result=export_result, plan=base.plan, - run_id=token.run_id, + run_id=base.run_id, state_details_sink=export_state_details, ) raw_worker_metadata = metrics.get("execution_workers", []) @@ -3810,7 +2359,9 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes else [] ) state_details = ( - _component_state_details(tuple(stage_evidence)) + _component_state_details( + tuple(stage_evidence), policy.execution_plan.final_stage_id + ) if policy.state_layout == "component" else export_state_details ) @@ -3836,7 +2387,11 @@ def execute(self, envelope: TorchRuntimeExecutionEnvelope) -> WorkerExecutionRes }, ) finally: - prepared.close() + try: + prepared.close() + finally: + for directory in reversed(owned_storage_dirs): + directory.cleanup() @DeveloperAPI diff --git a/src/tributo/integrations/sources/ray_torch.py b/src/tributo/integrations/sources/ray_torch.py index 81c9e42..5f3772d 100644 --- a/src/tributo/integrations/sources/ray_torch.py +++ b/src/tributo/integrations/sources/ray_torch.py @@ -167,9 +167,6 @@ def _open_source( raise AlgorithmConfigurationError( "Torch checkpoint implementation code digest drifted" ) - implementation_ref = QualifiedReference.parse(qualified) - _validate_module_digest(implementation_ref, options.implementation_code_digest) - implementation = _load_reference(implementation_ref) if options.policy_digest != descriptor.policy_digest: raise AlgorithmConfigurationError( "Torch export Policy digest does not match checkpoint" @@ -178,6 +175,13 @@ def _open_source( raise AlgorithmConfigurationError( "Torch export plan digest does not match checkpoint" ) + if descriptor.input_binding_digest != options.input_binding_digest: + raise AlgorithmConfigurationError( + "Torch export input binding digest does not match checkpoint" + ) + implementation_ref = QualifiedReference.parse(qualified) + _validate_module_digest(implementation_ref, options.implementation_code_digest) + implementation = _load_reference(implementation_ref) if not isinstance(implementation, type) or not issubclass( implementation, (TorchRecipe, RayTorchAdapter) ): @@ -209,7 +213,6 @@ def _open_source( state_layout=descriptor.state_layout, adapter_identity=descriptor.adapter_identity, resume_supported=descriptor.resume_supported, - same_world_size_resume=descriptor.same_world_size_resume, ) from tributo.algorithms.spi import TorchStageContext @@ -285,7 +288,6 @@ def _open_source( state_layout=descriptor.state_layout, adapter_identity=descriptor.adapter_identity, resume_supported=descriptor.resume_supported, - same_world_size_resume=descriptor.same_world_size_resume, ) from tributo.algorithms.spi import TorchStageContext @@ -321,25 +323,6 @@ def _read_descriptor(root: Path) -> TorchCheckpointDescriptor: payload = json.loads(path.read_text(encoding="utf-8")) descriptor = TorchCheckpointDescriptor.from_dict(payload) root_resolved = root.resolve() - commit_path = root / "torch_stage_commit.json" - if commit_path.exists() or commit_path.is_symlink(): - if commit_path.is_symlink() or not commit_path.is_file(): - raise AlgorithmConfigurationError( - "Torch checkpoint commit marker is invalid" - ) - try: - commit = json.loads(commit_path.read_text(encoding="utf-8")) - except (OSError, TypeError, ValueError) as exc: - raise AlgorithmConfigurationError( - "Torch checkpoint commit marker is malformed" - ) from exc - if not isinstance(commit, Mapping) or ( - commit.get("identity") != descriptor.identity.to_dict() - or commit.get("descriptor_digest") != descriptor.digest - ): - raise AlgorithmConfigurationError( - "Torch checkpoint commit marker does not match descriptor" - ) actual_files: dict[str, str] = {} for candidate in sorted(root.rglob("*")): if candidate.is_symlink() or not candidate.resolve().is_relative_to( @@ -350,7 +333,6 @@ def _read_descriptor(root: Path) -> TorchCheckpointDescriptor: ) if candidate.is_file() and candidate.name not in { "torch_checkpoint_descriptor.json", - "torch_stage_commit.json", ".metadata.json", }: actual_files[candidate.relative_to(root).as_posix()] = _sha256_file( @@ -403,7 +385,6 @@ def _load_recipe_model( state_layout=descriptor.state_layout, adapter_identity=descriptor.adapter_identity, resume_supported=descriptor.resume_supported, - same_world_size_resume=descriptor.same_world_size_resume, ) from tributo.algorithms.spi import ( TorchBuildContext, diff --git a/src/tributo/training/portable_tune.py b/src/tributo/training/portable_tune.py index 474ee5e..95a0b54 100644 --- a/src/tributo/training/portable_tune.py +++ b/src/tributo/training/portable_tune.py @@ -12,6 +12,7 @@ from tributo._common.immutable import deep_thaw from tributo.algorithms.api import ( DistributedAlgorithmDescriptor, + DistributionStrategy, ExecutionRequest, ResolvedAlgorithmPlan, ResultPolicy, @@ -88,6 +89,14 @@ def _trial_request( ) +def _trial_checkpoint_enabled(plan: ResolvedAlgorithmPlan) -> bool: + """Return whether the existing portable checkpoint files apply to this plan.""" + return not ( + plan.distribution_spec is not None + and plan.distribution_spec.strategy is DistributionStrategy.RAY_TRAIN_TORCH + ) + + @PublicAPI(stability="alpha") class PortableTuneRunner: """Tune one Wheel algorithm through its normal distributed Core Runtime.""" @@ -173,6 +182,9 @@ def trainable(sampled_values: dict[str, Any]) -> None: ) metric = _extract_target_metric(result.execution.metrics, metric_name) metrics = {metric_name: metric} + if not _trial_checkpoint_enabled(plan): + ray_tune.report(metrics) + return manifest = checkpoint_dir / "manifest.json" state = checkpoint_dir / "state.bin" if manifest.is_file() and state.is_file(): diff --git a/tests/algorithms/test_conformance_cli.py b/tests/algorithms/test_conformance_cli.py new file mode 100644 index 0000000..9e324ff --- /dev/null +++ b/tests/algorithms/test_conformance_cli.py @@ -0,0 +1,109 @@ +"""Installed-Wheel conformance CLI contracts.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import tributo.algorithms.conformance as conformance +import tributo.algorithms.conformance_cli as conformance_cli +from tributo.algorithms.conformance import AlgorithmPackageConformanceReport + + +def test_installed_conformance_matches_identity_and_contract_manifest( + monkeypatch, tmp_path +) -> None: + manifest = tmp_path / "identities.json" + manifest.write_text( + json.dumps( + { + "schema_version": 1, + "entry_points": { + "example": { + "algorithm_id": "example", + "distribution": "tributo-algorithms-example", + "implementation_id": "example.implementation", + } + }, + } + ), + encoding="utf-8", + ) + descriptor = SimpleNamespace( + registration=SimpleNamespace( + contract_bindings=SimpleNamespace( + config=object(), input=object(), output=object(), coverage=object() + ) + ) + ) + + class Distribution: + metadata = {"Name": "tributo-algorithms-example"} + + class EntryPoint: + name = "example" + dist = Distribution() + + @staticmethod + def load(): + return descriptor + + monkeypatch.setattr( + conformance.importlib.metadata, + "entry_points", + lambda **kwargs: (EntryPoint(),), + ) + monkeypatch.setattr( + conformance, + "validate_installed_algorithm_package", + lambda descriptor, entry_point_name: AlgorithmPackageConformanceReport( + algorithm_id="example", + implementation_id="example.implementation", + distribution="tributo-algorithms-example", + package_version="1.0.0", + entry_point_name=entry_point_name, + contract_ids=("config", "input", "output", "coverage"), + ), + ) + + reports = conformance._run_installed_conformance( + distribution_prefix="tributo-algorithms-", + expected_count=1, + identity_manifest=manifest, + required_contracts=("config", "input", "output", "coverage"), + forbidden_imports=(), + ) + + assert len(reports) == 1 + assert reports[0].entry_point_name == "example" + + +def test_conformance_cli_prints_deterministic_report(monkeypatch, capsys) -> None: + report = AlgorithmPackageConformanceReport( + algorithm_id="example", + implementation_id="example.implementation", + distribution="tributo-algorithms-example", + package_version="1.0.0", + entry_point_name="example", + contract_ids=("config", "input", "output", "coverage"), + ) + monkeypatch.setattr( + conformance_cli, + "_run_installed_conformance", + lambda **kwargs: (report,), + ) + + assert ( + conformance_cli.main( + [ + "--distribution-prefix", + "tributo-algorithms-", + "--expected-count", + "1", + "--identity-manifest", + "identities.json", + ] + ) + == 0 + ) + assert json.loads(capsys.readouterr().out)["count"] == 1 diff --git a/tests/algorithms/test_execution_contract.py b/tests/algorithms/test_execution_contract.py index e9f04ff..43d0a51 100644 --- a/tests/algorithms/test_execution_contract.py +++ b/tests/algorithms/test_execution_contract.py @@ -14,6 +14,7 @@ AlgorithmConfigurationError, AlgorithmExecutionResult, AlgorithmOperation, + ComponentStageEvidence, DistributionStrategy, ExecutionProfile, ExecutionReceipt, @@ -110,6 +111,22 @@ def test_worker_and_state_evidence_reject_truthy_string_coercion() -> None: StateCoordinationEvidence.from_dict(state) +@pytest.mark.parametrize("worker_digest", [None, "c" * 64]) +def test_component_stage_rejects_unbound_worker_state_digest( + worker_digest: str | None, +) -> None: + with pytest.raises(AlgorithmConfigurationError, match="model digests"): + ComponentStageEvidence( + stage_id="teacher", + workers=( + _worker(0), + replace(_worker(1), model_state_digest=worker_digest), + ), + roles=(), + state_digest="a" * 64, + ) + + def test_local_multi_worker_proves_model_distribution_not_cross_node() -> None: receipt = _receipt(ExecutionProfile.LOCAL) @@ -121,6 +138,7 @@ def test_local_multi_worker_proves_model_distribution_not_cross_node() -> None: assert receipt.to_dict()["distributed"] is True assert receipt.to_dict()["runtime_owned"] is False assert receipt.to_dict()["resource_preflight"] == "validated" + assert "torch_evidence" not in receipt.to_dict() def test_coordinator_receipt_preserves_requested_alias_and_canonical_algorithm() -> ( diff --git a/tests/algorithms/test_support_evidence.py b/tests/algorithms/test_support_evidence.py index adc9832..47de997 100644 --- a/tests/algorithms/test_support_evidence.py +++ b/tests/algorithms/test_support_evidence.py @@ -2,6 +2,8 @@ from __future__ import annotations +import hashlib +import json from dataclasses import replace from datetime import datetime, timedelta, timezone @@ -140,6 +142,34 @@ def test_descriptor_cannot_self_grant_tested_or_supported() -> None: assert record.validated_execution_profiles == () +def test_non_torch_evidence_id_keeps_the_pre_torch_payload() -> None: + descriptor = _descriptor() + issued_at = datetime(2026, 1, 1, tzinfo=timezone.utc) + evidence = _evidence(descriptor, issued_at=issued_at) + expected = { + "algorithm_id": evidence.algorithm_id, + "implementation_id": evidence.implementation_id, + "distribution": evidence.distribution, + "package_version": evidence.package_version, + "wheel_sha256": evidence.wheel_sha256, + "descriptor_api_version": evidence.descriptor_api_version, + "contract_digests": evidence.contract_digests, + "distributed_semantics": evidence.distributed_semantics.value, + "execution_profile": evidence.execution_profile.value, + "issuer": evidence.issuer, + "issued_at": issued_at.isoformat(), + "gate": evidence.gate, + "result_reference": evidence.result_reference, + } + + assert ( + evidence.evidence_id + == hashlib.sha256( + json.dumps(expected, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + ) + + def test_exact_trusted_wheel_evidence_grants_support() -> None: descriptor = _descriptor() overlay = AlgorithmSupportEvidenceRegistry( diff --git a/tests/algorithms/test_torch_recipe_worker.py b/tests/algorithms/test_torch_recipe_worker.py index bce4ff8..eb35cdf 100644 --- a/tests/algorithms/test_torch_recipe_worker.py +++ b/tests/algorithms/test_torch_recipe_worker.py @@ -2,7 +2,6 @@ from __future__ import annotations -import json from pathlib import Path from typing import Any @@ -13,6 +12,7 @@ from tributo.algorithms import ( TorchBatch, TorchLossContribution, + TorchMetricContribution, TorchMetricPlan, TorchModuleSet, TorchOptimizationPlan, @@ -22,11 +22,13 @@ TorchStageRunIdentity, TorchStepResult, ) +from tributo.algorithms.api import AlgorithmExecutionError from tributo.integrations.algorithm_runtimes.ray_train_torch import ( torch_recipe_train_loop_per_worker, ) torch = pytest.importorskip("torch") +ray_train_torch = pytest.importorskip("ray.train.torch") class BinaryLinearRecipe(TorchRecipe): @@ -79,6 +81,18 @@ def artifact_plan(self, context: object) -> dict[str, object]: return {"source_kind": "torch_module"} +class UnexpectedMetricRecipe(BinaryLinearRecipe): + def training_step( + self, modules: TorchModuleSet, batch: TorchBatch, context: object + ) -> TorchStepResult: + step = super().training_step(modules, batch, context) + return TorchStepResult( + outputs=step.outputs, + loss=step.loss, + metrics={"unexpected": TorchMetricContribution(1.0, 1.0)}, + ) + + class _Iterator: def __init__(self, batches: list[dict[str, Any]]) -> None: self._batches = batches @@ -126,6 +140,15 @@ def test_recipe_worker_reports_typed_checkpoint_and_exact_coverage( }, ] ) + val_data = _Iterator( + [ + { + "x1": torch.tensor([0.5]), + "x2": torch.tensor([0.5]), + "label": torch.tensor([1.0]), + } + ] + ) reports: list[tuple[dict[str, Any], set[str]]] = [] def report(metrics: dict[str, Any], checkpoint: Any | None = None) -> None: @@ -135,16 +158,30 @@ def report(metrics: dict[str, Any], checkpoint: Any | None = None) -> None: (dict(metrics), {path.name for path in Path(directory).iterdir()}) ) - monkeypatch.setattr(ray.train.torch, "prepare_model", lambda model: model) + monkeypatch.setattr(ray_train_torch, "prepare_model", lambda model: model) monkeypatch.setattr(ray.train, "get_context", lambda: _TrainContext()) monkeypatch.setattr( ray.train, "get_dataset_shard", - lambda name: data if name == "train" else (_ for _ in ()).throw(KeyError(name)), + lambda name: ( + data + if name == "train" + else val_data + if name == "val" + else (_ for _ in ()).throw(KeyError(name)) + ), + ) + monkeypatch.setattr( + ray.train, + "get_checkpoint", + lambda: (_ for _ in ()).throw(AssertionError("retry state must not be read")), ) - monkeypatch.setattr(ray.train, "get_checkpoint", lambda: None) monkeypatch.setattr(ray.train, "report", report) monkeypatch.setattr(ray, "get_runtime_context", lambda: _RuntimeContext()) + monkeypatch.setattr( + "tributo.integrations.algorithm_runtimes.ray_train_torch._validate_module_digest", + lambda reference, digest: None, + ) identity = TorchStageRunIdentity( "aabbccdd", "11223344", @@ -165,7 +202,7 @@ def report(metrics: dict[str, Any], checkpoint: Any | None = None) -> None: identity, input_binding_digest="3" * 64, ) - stage = TorchStageContext(runtime, "train", 0, True, ("train",)) + stage = TorchStageContext(runtime, "train", 0, True, ("train", "val")) torch_recipe_train_loop_per_worker( { "training": {"epochs": 1, "batch_size": 2}, @@ -176,60 +213,20 @@ def report(metrics: dict[str, Any], checkpoint: Any | None = None) -> None: } ) assert reports and reports[0][0]["checkpoint_descriptor"]["completed_step"] == 2 + assert reports[0][0]["execution_workers"][0]["input_rows"] == { + "train": 3, + "val": 1, + } assert { "model.pt", - "optimizer.pt", - "scaler.pt", - "rng_state.pt", + "torch_execution_evidence.json", "torch_checkpoint_descriptor.json", } <= reports[0][1] -def test_recipe_worker_rejects_stale_retry_checkpoint( +def test_recipe_worker_replays_ray_retry_from_stage_start( monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, ) -> None: - from tributo.algorithms import TorchCheckpointDescriptor - - recipe = BinaryLinearRecipe() - modules = recipe.build_modules(None) - model = modules["model"] - optimizer = recipe.configure_optimizers(modules, None).optimizer - torch.save(model.state_dict(), tmp_path / "model.pt") - torch.save(optimizer.state_dict(), tmp_path / "optimizer.pt") - torch.save({}, tmp_path / "scaler.pt") - (tmp_path / "rng_state.pt").write_bytes(torch.get_rng_state().numpy().tobytes()) - identity = TorchStageRunIdentity( - "aabbccdd", - "11223344", - "train", - 1, - "example", - "example.binary", - "0" * 64, - "1" * 64, - "2" * 64, - ) - payload_files = { - name: __import__("hashlib").sha256((tmp_path / name).read_bytes()).hexdigest() - for name in ("model.pt", "optimizer.pt", "scaler.pt", "rng_state.pt") - } - descriptor = TorchCheckpointDescriptor( - 1, - identity, - identity.run_config_name, - "replicated", - 1, - 1, - "1" * 64, - "2" * 64, - "3" * 64, - "0" * 64, - payload_files, - ) - (tmp_path / "torch_checkpoint_descriptor.json").write_text( - json.dumps(descriptor.to_dict()), encoding="utf-8" - ) data = _Iterator( [ { @@ -248,75 +245,132 @@ def get_dataset_shard(name: str) -> _Iterator: monkeypatch.setattr(ray.train, "get_context", lambda: _TrainContext()) monkeypatch.setattr(ray.train, "get_dataset_shard", get_dataset_shard) - monkeypatch.setattr(ray.train.torch, "prepare_model", lambda model: model) + monkeypatch.setattr(ray_train_torch, "prepare_model", lambda model: model) - class _Checkpoint: - def as_directory(self): - from contextlib import contextmanager - - @contextmanager - def opened(): - yield tmp_path - - return opened() - - monkeypatch.setattr(ray.train, "get_checkpoint", lambda: _Checkpoint()) + monkeypatch.setattr( + ray.train, + "get_checkpoint", + lambda: (_ for _ in ()).throw(AssertionError("retry state must not be read")), + ) monkeypatch.setattr( ray.train, "report", lambda metrics, checkpoint=None: reports.append(dict(metrics)), ) monkeypatch.setattr(ray, "get_runtime_context", lambda: _RuntimeContext()) + monkeypatch.setattr( + "tributo.integrations.algorithm_runtimes.ray_train_torch._validate_module_digest", + lambda reference, digest: None, + ) - torch_recipe_train_loop_per_worker( - { - "model": {"input_features": 2}, - "optimizer": {"learning_rate": 0.1}, - "training": {"epochs": 2, "batch_size": 2, "seed": 7}, - "ray": {"resume": {"checkpoint_interval": 1}}, - "_tributo_recipe_ref": ("tests.support.torch_recipe:BinaryLinearRecipe"), - "_tributo_recipe_code_digest": None, - "_tributo_implementation_id": "example.binary_linear", - "_tributo_algorithm": "binary_linear", - "_tributo_feature_names": ["x1", "x2"], - "_tributo_label_name": "label", - "_tributo_weight_name": None, - "_tributo_input_binding_digest": "a" * 64, - "_tributo_distribution_spec_digest": "b" * 64, - "_tributo_resume_from": str(tmp_path), - "_tributo_metric_reducers": { - "accuracy": "sum_count", - "train_loss": "sum_count", - }, - "_core_implementation_ref": "tests.support.torch_recipe:BinaryLinearRecipe", - "_core_implementation_code_digest": "0" * 64, - "_core_input_binding_digest": "a" * 64, - "_core_stage_context": TorchStageContext( - TorchRuntimeContext( - {}, - "example.binary_linear", + config = { + "model": {"input_features": 2}, + "optimizer": {"learning_rate": 0.1}, + "training": {"epochs": 2, "batch_size": 2, "seed": 7}, + "_core_implementation_ref": "tests.support.torch_recipe:BinaryLinearRecipe", + "_core_implementation_code_digest": "0" * 64, + "_core_input_binding_digest": "a" * 64, + "_core_feature_names": ["x1", "x2"], + "_core_label_name": "label", + "_core_stage_context": TorchStageContext( + TorchRuntimeContext( + {}, + "example.binary_linear", + 1, + "1" * 64, + "2" * 64, + TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", 1, + "binary", + "example.binary_linear", + "0" * 64, "1" * 64, "2" * 64, - TorchStageRunIdentity( - "aabbccdd", - "11223344", - "train", - 1, - "binary", - "example.binary_linear", - "0" * 64, - "1" * 64, - "2" * 64, - ), - input_binding_digest="a" * 64, ), - "train", - 0, - True, - ("train",), - ).to_dict(), - }, + input_binding_digest="a" * 64, + ), + "train", + 0, + True, + ("train",), + ).to_dict(), + } + torch.manual_seed(101) + torch_recipe_train_loop_per_worker(config) + torch.manual_seed(202) + torch_recipe_train_loop_per_worker(config) + + assert len(reports) == 2 + assert all( + report["checkpoint_descriptor"]["completed_step"] == 2 for report in reports ) + assert reports[0]["model_state_digest"] == reports[1]["model_state_digest"] - assert reports and reports[0]["checkpoint_descriptor"]["completed_step"] == 3 + +def test_recipe_worker_rejects_metric_plan_drift( + monkeypatch: pytest.MonkeyPatch, +) -> None: + data = _Iterator( + [ + { + "x1": torch.tensor([0.0]), + "x2": torch.tensor([0.0]), + "label": torch.tensor([0.0]), + } + ] + ) + monkeypatch.setattr(ray.train, "get_context", lambda: _TrainContext()) + monkeypatch.setattr( + ray.train, + "get_dataset_shard", + lambda name: data if name == "train" else (_ for _ in ()).throw(KeyError(name)), + ) + monkeypatch.setattr(ray_train_torch, "prepare_model", lambda model: model) + monkeypatch.setattr(ray, "get_runtime_context", lambda: _RuntimeContext()) + monkeypatch.setattr( + "tributo.integrations.algorithm_runtimes.ray_train_torch._validate_module_digest", + lambda reference, digest: None, + ) + stage = TorchStageContext( + TorchRuntimeContext( + {}, + "example.unexpected_metric", + 1, + "1" * 64, + "2" * 64, + TorchStageRunIdentity( + "aabbccdd", + "11223344", + "train", + 1, + "example", + "example.unexpected_metric", + "0" * 64, + "1" * 64, + "2" * 64, + ), + input_binding_digest="3" * 64, + ), + "train", + 0, + True, + ("train",), + ) + with pytest.raises( + AlgorithmExecutionError, + match="metrics do not match TorchMetricPlan", + ): + torch_recipe_train_loop_per_worker( + { + "training": {"epochs": 1, "batch_size": 1}, + "_core_implementation_ref": ( + "tests.algorithms.test_torch_recipe_worker:UnexpectedMetricRecipe" + ), + "_core_implementation_code_digest": "0" * 64, + "_core_input_binding_digest": "3" * 64, + "_core_stage_context": stage.to_dict(), + } + ) diff --git a/tests/algorithms/test_torch_runtime_contract_v1.py b/tests/algorithms/test_torch_runtime_contract_v1.py index da3041a..5074e5e 100644 --- a/tests/algorithms/test_torch_runtime_contract_v1.py +++ b/tests/algorithms/test_torch_runtime_contract_v1.py @@ -19,8 +19,7 @@ TorchAccumulationWindow, TorchBackwardContext, TorchCheckpointDescriptor, - TorchCheckpointLocator, - TorchCheckpointProgress, + TorchCheckpointPayloadDraft, TorchCompositeLossContribution, TorchDatasetRoute, TorchGlobalLossReduction, @@ -29,18 +28,17 @@ TorchMetricPolicy, TorchMetricReductionContext, TorchPolicy, - TorchPreflightLease, - TorchPreflightTokenData, - TorchRankProgressStatistics, - TorchRecoveryEnvelope, TorchStageRunIdentity, TorchStageSpec, apply_torch_loss_backward, reduce_torch_metrics, report_torch_checkpoint, - torch_run_config_name, ) -from tributo.algorithms.spi import TorchRuntimeContext, TorchStageContext +from tributo.algorithms.spi import ( + TorchOptimizationPlan, + TorchRuntimeContext, + TorchStageContext, +) class _Scalar: @@ -56,25 +54,9 @@ def item(self) -> float: return self.value -def _identity_kwargs() -> dict[str, object]: - return { - "run_id": "aabbccdd", - "invocation_id": "11223344", - "algorithm": "example", - "implementation_ref": "example:Recipe", - "implementation_code_digest": "0" * 64, - "policy_digest": "1" * 64, - "execution_plan_digest": "2" * 64, - "runtime_id": "tributo.ray_train_torch", - "plan_digest": "3" * 64, - } - - def test_torch_policy_and_run_name_are_deterministic() -> None: route = TorchDatasetRoute("train", "split_exact") - execution_plan = SingleStageTorchPlan( - stage=TorchStageSpec("train", "example:loop", ("train",)) - ) + execution_plan = SingleStageTorchPlan(stage=TorchStageSpec("train", ("train",))) policy = TorchPolicy( torch_runtime_api_version=1, loop_owner="core_recipe", @@ -96,30 +78,7 @@ def test_torch_policy_and_run_name_are_deterministic() -> None: policy.digest, execution_plan.digest, ) - assert torch_run_config_name(identity) == identity.run_config_name - - -def test_preflight_lease_is_one_shot_and_identity_bound() -> None: - lease = TorchPreflightLease(TorchPreflightTokenData(**_identity_kwargs())) - lease.claim( - run_id="aabbccdd", - invocation_id="11223344", - plan_digest="3" * 64, - runtime_id="tributo.ray_train_torch", - ) - lease.consume( - run_id="aabbccdd", - invocation_id="11223344", - plan_digest="3" * 64, - runtime_id="tributo.ray_train_torch", - ) - with pytest.raises(AlgorithmExecutionError): - lease.consume( - run_id="aabbccdd", - invocation_id="11223344", - plan_digest="3" * 64, - runtime_id="tributo.ray_train_torch", - ) + assert identity.run_config_name.startswith("tributo-torch-v1-") def test_loss_contribution_requires_zero_dimensional_scalar() -> None: @@ -134,6 +93,35 @@ def test_loss_contribution_requires_zero_dimensional_scalar() -> None: assert set(composite.differentiable_components) == {"loss_a", "loss_b", "loss_c"} +@pytest.mark.parametrize( + "max_gradient_norm", + [True, 0, -1, float("nan"), float("inf"), float("-inf")], +) +def test_torch_optimization_plan_rejects_invalid_gradient_clip_norm( + max_gradient_norm: float | int | bool, +) -> None: + with pytest.raises(ValueError, match="finite and positive"): + TorchOptimizationPlan( + optimizer=object(), + max_gradient_norm=max_gradient_norm, + ) + + +def test_torch_optimization_plan_accepts_optional_positive_gradient_clip_norm() -> None: + assert ( + TorchOptimizationPlan( + optimizer=object(), max_gradient_norm=None + ).max_gradient_norm + is None + ) + assert ( + TorchOptimizationPlan( + optimizer=object(), max_gradient_norm=1.0 + ).max_gradient_norm + == 1.0 + ) + + def test_backward_and_metric_helpers_use_explicit_normalizers() -> None: events: list[object] = [] result = apply_torch_loss_backward( @@ -184,15 +172,11 @@ def test_backward_helper_reduces_only_the_accumulation_window_total() -> None: assert scales == [2 / 10] -def test_locator_rejects_local_paths_and_policy_replicate_budget_is_explicit() -> None: - with pytest.raises(AlgorithmConfigurationError): - TorchCheckpointLocator("/tmp/checkpoint", "0" * 64) +def test_policy_replicate_budget_is_explicit() -> None: route = TorchDatasetRoute( "nodes", "replicate", max_rows=10, max_bytes_per_worker=10 ) - plan = SingleStageTorchPlan( - stage=TorchStageSpec("train", "example:loop", ("nodes",)) - ) + plan = SingleStageTorchPlan(stage=TorchStageSpec("train", ("nodes",))) with pytest.raises(AlgorithmConfigurationError): TorchPolicy( 1, @@ -205,6 +189,22 @@ def test_locator_rejects_local_paths_and_policy_replicate_budget_is_explicit() - ) +def test_torch_v1_rejects_cross_run_recovery() -> None: + route = TorchDatasetRoute("train", "split_exact") + plan = SingleStageTorchPlan(stage=TorchStageSpec("train", ("train",))) + with pytest.raises(AlgorithmConfigurationError, match="cross-Run"): + TorchPolicy( + 1, + "core_recipe", + "torch.ddp.replicated", + (route,), + plan, + "replicated", + {"train_loss": MetricReduction.SUM_COUNT}, + resume_supported=True, + ) + + def test_composite_global_state_keeps_component_and_normalizer_names_independent() -> ( None ): @@ -218,35 +218,7 @@ def test_composite_global_state_keeps_component_and_normalizer_names_independent assert set(state.normalizers) == {"positive_count", "negative_count"} -def test_stage_dependency_is_allowed_when_external_recovery_is_disabled() -> None: - from tributo.integrations.algorithm_runtimes.ray_train_torch import ( - _control_for_stage, - ) - - control = _control_for_stage( - SimpleNamespace( - algorithm_config={}, - runtime=SimpleNamespace(resume_from=None), - ), - SimpleNamespace( - resume_supported=False, - digest="1" * 64, - execution_plan=SimpleNamespace(digest="2" * 64), - ), - SimpleNamespace(stage_id="student", checkpoint_from_stage="teacher"), - run_id="aabbccdd", - invocation_id="11223344", - predecessor={ - "locator": "s3://bucket/teacher-checkpoint", - "descriptor_digest": "3" * 64, - }, - ) - assert control is not None - assert control["purpose"] == "stage_dependency" - assert control["source_stage_id"] == "teacher" - - -def test_role_evidence_falls_back_to_primary_binding_for_alias_roles() -> None: +def test_role_evidence_requires_an_actual_role_binding() -> None: from tributo.integrations.algorithm_runtimes.ray_train_torch import ( _binding_digest_for_role, ) @@ -255,11 +227,11 @@ class Descriptors: def get(self, role: str) -> object: raise AlgorithmConfigurationError(f"unknown resolved input role: {role}") - primary = SimpleNamespace(binding_digest="3" * 64) plan = SimpleNamespace( - input_descriptors=Descriptors(), primary_input_descriptor=primary + input_descriptors=Descriptors(), ) - assert _binding_digest_for_role(plan, "val") == "3" * 64 + with pytest.raises(AlgorithmConfigurationError, match="no binding digest"): + _binding_digest_for_role(plan, "val") def test_worker_evidence_defaults_only_missing_declared_resources() -> None: @@ -310,15 +282,42 @@ def test_component_state_details_project_stage_coverage() -> None: to_dict=lambda: {"stage_id": "finetune", "state": "b" * 64}, ), ) - details = _component_state_details(stages) + details = _component_state_details(stages, "pretrain") assert details["component_stage_count"] == 2 assert details["component_stages"] == "pretrain,finetune" - assert details["anchor_stage"] == "finetune" + assert details["anchor_stage"] == "pretrain" assert details["stage.pretrain.rows"] == 16 assert details["stage.finetune.rows"] == 12 assert len(details["composition_digest"]) == 64 +def test_component_stage_metrics_keep_named_metrics_and_final_train_loss() -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _record_stage_metrics, + ) + + metrics: dict[str, float] = {} + declared = {"train_loss", "teacher_loss", "student_loss"} + _record_stage_metrics( + metrics, + {"train_loss": 2.0, "teacher_loss": 2.0}, + declared, + is_final=False, + ) + _record_stage_metrics( + metrics, + {"train_loss": 1.0, "student_loss": 1.0}, + declared, + is_final=True, + ) + + assert metrics == { + "train_loss": 1.0, + "teacher_loss": 2.0, + "student_loss": 1.0, + } + + def test_source_state_details_preserve_adapter_declared_scalars() -> None: from tributo.integrations.algorithm_runtimes.ray_train_torch import ( _source_state_details, @@ -339,6 +338,65 @@ def test_source_state_details_preserve_adapter_declared_scalars() -> None: } +def test_torch_ray_config_rejects_removed_resume_options() -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _torch_ray_config, + ) + + plan = SimpleNamespace( + algorithm_config={"ray": {"storage_path": "/tmp/ray", "max_failures": 1}} + ) + assert _torch_ray_config(plan) == ("/tmp/ray", 1) + plan.algorithm_config["ray"]["resume"] = {"checkpoint_interval": 1} + with pytest.raises(AlgorithmConfigurationError, match="unsupported key"): + _torch_ray_config(plan) + + +def test_runtime_rejects_reserved_torch_policy_features() -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _validate_runtime_policy, + ) + + supported = SimpleNamespace( + state_layout="replicated", + dataset_routing=(SimpleNamespace(mode="split_exact"),), + checkpoint_owner_rank=0, + checkpoint_adapter_ref=None, + evidence_adapter_ref=None, + ) + _validate_runtime_policy(supported) + for policy, message in ( + ( + SimpleNamespace(**{**vars(supported), "state_layout": "sharded"}), + "sharded", + ), + ( + SimpleNamespace(**{**vars(supported), "checkpoint_owner_rank": 1}), + "checkpoint_owner_rank", + ), + ( + SimpleNamespace( + **{ + **vars(supported), + "dataset_routing": (SimpleNamespace(mode="split_framework"),), + } + ), + "split_framework", + ), + ( + SimpleNamespace( + **{ + **vars(supported), + "checkpoint_adapter_ref": "example:checkpoint", + } + ), + "adapter references", + ), + ): + with pytest.raises(AlgorithmConfigurationError, match=message): + _validate_runtime_policy(policy) + + def test_replicated_role_evidence_uses_per_rank_rows() -> None: from tributo.algorithms.api import TorchRoleExecutionEvidence @@ -351,6 +409,7 @@ def test_replicated_role_evidence_uses_per_rank_rows() -> None: expected_rows=8, observed_rows=8, rows_per_rank=(8, 8), + replicated_bytes_per_worker=123, ) assert evidence.rows_per_rank == (8, 8) with pytest.raises(AlgorithmConfigurationError): @@ -363,9 +422,58 @@ def test_replicated_role_evidence_uses_per_rank_rows() -> None: expected_rows=8, observed_rows=8, rows_per_rank=(), + replicated_bytes_per_worker=123, ) +def test_stage_route_validation_returns_actual_replicated_bytes() -> None: + from tributo.integrations.algorithm_runtimes.ray_train_torch import ( + _validate_stage_routes, + ) + + calls: list[str] = [] + + class MaterializedDataset: + def count(self) -> int: + calls.append("count") + return 8 + + def size_bytes(self) -> int: + calls.append("size_bytes") + return 123 + + materialized = MaterializedDataset() + + class LimitedDataset: + def materialize(self) -> MaterializedDataset: + calls.append("materialize") + return materialized + + class Dataset: + def count(self) -> int: + raise AssertionError("the unbounded Dataset must not be counted") + + def limit(self, count: int) -> LimitedDataset: + calls.append(f"limit:{count}") + assert count == 9 + return LimitedDataset() + + route = TorchDatasetRoute( + "nodes", "replicate", max_rows=8, max_bytes_per_worker=256 + ) + datasets: dict[str, object] = {"nodes": Dataset()} + rows, replicated_bytes = _validate_stage_routes( + SimpleNamespace(dataset_routing=(route,), max_replicated_bytes_per_worker=256), + SimpleNamespace(input_roles=("nodes",)), + datasets, + 2, + ) + assert rows == {"nodes": 8} + assert replicated_bytes == {"nodes": 123} + assert datasets["nodes"] is materialized + assert calls == ["limit:9", "materialize", "count", "size_bytes"] + + def test_adapter_worker_config_cannot_carry_core_paths() -> None: from tributo.integrations.algorithm_runtimes.ray_train_torch import ( _validate_adapter_worker_config, @@ -536,7 +644,9 @@ def export_bundle(self, source, config, *, tributo_version): assert execution.outputs["composition_digest"] == "a" * 64 -def test_checkpoint_report_builds_descriptor_from_payload(tmp_path) -> None: +def test_checkpoint_report_builds_descriptor_from_payload( + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: identity = TorchStageRunIdentity( "aabbccdd", "11223344", @@ -561,18 +671,18 @@ def test_checkpoint_report_builds_descriptor_from_payload(tmp_path) -> None: payload = tmp_path / "model.pt" payload.write_bytes(b"model") captured: dict[str, object] = {} - - class Draft: - checkpoint_dir = tmp_path - - def report(self, *, metrics, stage_context, completed_step) -> None: - captured.update(metrics) - assert stage_context is stage - assert completed_step == 1 + monkeypatch.setattr( + "ray.train.get_context", + lambda: SimpleNamespace(get_world_rank=lambda: 0), + ) + monkeypatch.setattr( + "ray.train.report", + lambda metrics, checkpoint=None: captured.update(metrics), + ) report_torch_checkpoint( {"train_loss": 0.5}, - Draft(), + TorchCheckpointPayloadDraft(tmp_path), stage, 1, ) @@ -584,34 +694,6 @@ def report(self, *, metrics, stage_context, completed_step) -> None: assert (tmp_path / "torch_checkpoint_descriptor.json").is_file() -def test_recovery_envelope_roundtrip_and_locator_digest_binding() -> None: - locator = TorchCheckpointLocator("s3://bucket/stage", "4" * 64) - envelope = TorchRecoveryEnvelope( - completed_stage_ids=("pretrain",), - stage_checkpoints={"pretrain": locator}, - active_stage_id="finetune", - active_checkpoint=TorchCheckpointLocator("s3://bucket/active", "5" * 64), - ) - restored = TorchRecoveryEnvelope.from_dict(envelope.to_dict()) - assert restored == envelope - with pytest.raises(AlgorithmConfigurationError): - TorchRecoveryEnvelope( - completed_stage_ids=("pretrain",), - stage_checkpoints={ - "pretrain": TorchCheckpointLocator("s3://bucket/stage", "4" * 64) - }, - active_stage_id="pretrain", - active_checkpoint=TorchCheckpointLocator("s3://bucket/active", "5" * 64), - ) - with pytest.raises(AlgorithmConfigurationError): - TorchRecoveryEnvelope.from_dict( - { - "completed_stage_ids": ["pretrain"], - "stage_checkpoints": {"pretrain": "not-a-locator"}, - } - ) - - def test_checkpoint_payload_rejects_symlinked_descriptor(tmp_path) -> None: identity = TorchStageRunIdentity( "aabbccdd", @@ -639,14 +721,8 @@ def test_checkpoint_payload_rejects_symlinked_descriptor(tmp_path) -> None: target.write_text("{}", encoding="utf-8") (tmp_path / "torch_checkpoint_descriptor.json").symlink_to(target) - class Draft: - checkpoint_dir = tmp_path - - def report(self, *, metrics, stage_context, completed_step) -> None: - del metrics, stage_context, completed_step - with pytest.raises(AlgorithmExecutionError): - report_torch_checkpoint({}, Draft(), stage, 1) + report_torch_checkpoint({}, TorchCheckpointPayloadDraft(tmp_path), stage, 1) def test_checkpoint_report_rejects_core_metadata_fields(tmp_path) -> None: @@ -673,29 +749,16 @@ def test_checkpoint_report_rejects_core_metadata_fields(tmp_path) -> None: stage = TorchStageContext(runtime, "train", 0, True, ("train",)) (tmp_path / "model.pt").write_bytes(b"model") - class Draft: - checkpoint_dir = tmp_path - - def report(self, *, metrics, stage_context, completed_step) -> None: - del metrics, stage_context, completed_step - with pytest.raises(AlgorithmConfigurationError): report_torch_checkpoint( - {"checkpoint_locator": "s3://bucket/private"}, Draft(), stage, 1 + {"checkpoint_locator": "s3://bucket/private"}, + TorchCheckpointPayloadDraft(tmp_path), + stage, + 1, ) -def test_checkpoint_progress_roundtrip_and_conditional_resume_serialization() -> None: - progress = TorchCheckpointProgress( - epoch=2, - micro_batch_cursor=3, - optimizer_step=7, - scheduler_step=2, - accumulation_steps=4, - dataset_cursor_by_rank={"0": 3, "1": 3}, - shuffle_seed=44, - ) - assert TorchCheckpointProgress.from_dict(progress.to_dict()) == progress +def test_checkpoint_descriptor_omits_unsupported_resume_state() -> None: identity = TorchStageRunIdentity( "aabbccdd", "11223344", @@ -720,25 +783,12 @@ def test_checkpoint_progress_roundtrip_and_conditional_resume_serialization() -> implementation_code_digest=identity.implementation_code_digest, payload_files={"model.pt": "4" * 64}, resume_supported=False, - same_world_size_resume=None, ) assert "same_world_size_resume" not in descriptor.to_dict() assert TorchCheckpointDescriptor.from_dict(descriptor.to_dict()) == descriptor -def test_rank_progress_statistics_and_runtime_context_are_typed() -> None: - statistics = TorchRankProgressStatistics( - rows_processed=4, - coverage_totals={"coverage.positive": 2}, - loss_numerator_total=3.0, - loss_normalizer_total=4.0, - metric_totals={"accuracy": (2.0, 4.0)}, - reducer_observation={"branch": "nnpu_normal"}, - ) - assert TorchRankProgressStatistics.from_dict(statistics.to_dict()) == statistics - with pytest.raises(AlgorithmConfigurationError): - TorchRankProgressStatistics.from_dict({"rows_processed": "four"}) - +def test_runtime_context_omits_unsupported_resume_state() -> None: runtime = TorchRuntimeContext( algorithm_config={}, implementation_id="example.adapter", @@ -746,7 +796,6 @@ def test_rank_progress_statistics_and_runtime_context_are_typed() -> None: policy_digest="1" * 64, execution_plan_digest="2" * 64, resume_supported=False, - same_world_size_resume=None, ) payload = runtime.to_dict() assert "same_world_size_resume" not in payload @@ -759,95 +808,33 @@ def test_rank_progress_statistics_and_runtime_context_are_typed() -> None: "input_roles": ["train"], } ) - assert restored.runtime.same_world_size_resume is None + assert not hasattr(restored.runtime, "same_world_size_resume") -def test_scheduler_boundary_and_recovery_commit_are_fail_closed(tmp_path) -> None: +def test_worker_stage_context_omits_driver_output_paths() -> None: from tributo.integrations.algorithm_runtimes.ray_train_torch import ( - _require_checkpoint_commit, - _should_apply_epoch_scheduler, + _worker_stage_context, ) - assert _should_apply_epoch_scheduler( - restore_same_stage=True, - epoch=1, - restored_epoch=1, - restored_epoch_scheduler_applied=False, - ) - assert not _should_apply_epoch_scheduler( - restore_same_stage=True, - epoch=1, - restored_epoch=1, - restored_epoch_scheduler_applied=True, - ) - identity = TorchStageRunIdentity( - "aabbccdd", - "11223344", + stage = TorchStageContext( + TorchRuntimeContext( + algorithm_config={}, + implementation_id="example.adapter", + world_size=1, + policy_digest="1" * 64, + execution_plan_digest="2" * 64, + output_config={"bundle_uri": "/driver/model"}, + ), "train", - 1, - "example", - "example.recipe", - "0" * 64, - "1" * 64, - "2" * 64, + 0, + True, + ("train",), ) - (tmp_path / "model.pt").write_bytes(b"model") - descriptor = TorchCheckpointDescriptor( - schema_version=1, - identity=identity, - run_config_name=identity.run_config_name, - state_layout="replicated", - world_size=1, - completed_step=1, - policy_digest=identity.policy_digest, - execution_plan_digest=identity.execution_plan_digest, - input_binding_digest="3" * 64, - implementation_code_digest=identity.implementation_code_digest, - payload_files={"model.pt": "4" * 64}, - ) - with pytest.raises(AlgorithmExecutionError, match="commit"): - _require_checkpoint_commit(tmp_path, descriptor) + worker_context = _worker_stage_context(stage) -def test_local_stage_staging_ignores_prior_partial_attempt(tmp_path) -> None: - from tributo.integrations.algorithm_runtimes.ray_train_torch import ( - _persist_stage_checkpoint, - ) - - identity = TorchStageRunIdentity( - "aabbccdd", - "11223344", - "train", - 1, - "example", - "example.recipe", - "0" * 64, - "1" * 64, - "2" * 64, - ) - source = tmp_path / "source" - source.mkdir() - (source / "model.pt").write_bytes(b"model") - run_root = tmp_path / identity.run_config_name - run_root.mkdir() - - class Checkpoint: - @contextmanager - def as_directory(self): - yield source - - # A previous attempt with the same digest must not block a fresh staging - # attempt; only the committed destination is authoritative. - stale_path = run_root / f".stage_checkpoint.staging-{'4' * 64}-old" - stale_path.mkdir() - locator = _persist_stage_checkpoint( - Checkpoint(), - identity=identity, - storage_path=tmp_path, - descriptor_digest="4" * 64, - ) - assert locator == f"ray://{run_root / 'stage_checkpoint'}" - assert (run_root / "stage_checkpoint" / "torch_stage_commit.json").is_file() + assert stage.runtime.output_config == {"bundle_uri": "/driver/model"} + assert worker_context.runtime.output_config == {} def test_composite_zero_global_normalizer_fails_before_reducer( @@ -918,3 +905,18 @@ def test_removed_torch_public_surfaces_are_not_exported() -> None: assert not hasattr(algorithms, "TrainingRecipeV2") assert not hasattr(algorithms.AlgorithmBuilder, "from_torch_recipe") assert not hasattr(algorithms.AlgorithmBuilder, "from_training_recipe_v2") + for name in ( + "TorchCheckpointLocator", + "TorchCheckpointProgress", + "TorchPreflightLease", + "TorchPreflightTokenData", + "TorchRankProgressStatistics", + "TorchRecoveryEnvelope", + "TorchRuntimeExecutionEnvelope", + "TorchWorkerControlEnvelope", + "claim_torch_run_directory", + "describe_torch_checkpoint", + "torch_run_config_name", + "validate_torch_retry_identity", + ): + assert not hasattr(algorithms, name) diff --git a/tests/inference/test_contracts.py b/tests/inference/test_contracts.py index 36edda4..b78d871 100644 --- a/tests/inference/test_contracts.py +++ b/tests/inference/test_contracts.py @@ -2,7 +2,6 @@ from __future__ import annotations -import numpy as np import pytest from pydantic import ValidationError @@ -28,7 +27,6 @@ TensorInputBinding, TensorOutputBinding, ) -from tributo.inference.kernel import _build_input_tensor def _request(**updates) -> InferenceRequest: @@ -254,54 +252,6 @@ def test_single_column_mode_is_explicit_and_json_stable(self) -> None: TensorInputBinding.model_validate_json(scalar.model_dump_json()) == scalar ) - def test_single_vector_column_preserves_nested_tensor_rank(self) -> None: - column = np.empty(2, dtype=object) - column[0] = [[1.0], [2.0]] - column[1] = [[3.0], [4.0]] - tensor = _build_input_tensor( - {"window": column}, - columns=("window",), - dtype="float32", - single_column_mode="vector", - null_policy="error", - nan_policy="error", - ) - assert tensor.shape == (2, 2, 1) - - def test_multi_dimensional_object_column_preserves_nested_tensor_rank( - self, - ) -> None: - column = np.empty((2, 2), dtype=object) - column[0] = [np.asarray([1.0, 2.0]), np.asarray([3.0, 4.0])] - column[1] = [np.asarray([5.0, 6.0]), np.asarray([7.0, 8.0])] - tensor = _build_input_tensor( - {"window": column}, - columns=("window",), - dtype="float32", - single_column_mode="vector", - null_policy="error", - nan_policy="error", - ) - assert tensor.shape == (2, 2, 2) - - def test_arrow_nested_object_rows_preserve_tensor_rank(self) -> None: - column = np.empty(2, dtype=object) - column[0] = np.asarray( - [np.asarray([1.0, 2.0]), np.asarray([3.0, 4.0])], dtype=object - ) - column[1] = np.asarray( - [np.asarray([5.0, 6.0]), np.asarray([7.0, 8.0])], dtype=object - ) - tensor = _build_input_tensor( - {"window": column}, - columns=("window",), - dtype="float32", - single_column_mode="vector", - null_policy="error", - nan_policy="error", - ) - assert tensor.shape == (2, 2, 2) - def test_scalar_single_column_mode_rejects_invalid_contracts(self) -> None: with pytest.raises(ValidationError, match="requires exactly one column"): TensorInputBinding( diff --git a/tests/integration/test_distributed_algorithm_it_contract.py b/tests/integration/test_distributed_algorithm_it_contract.py index fa8323e..bfd0c51 100644 --- a/tests/integration/test_distributed_algorithm_it_contract.py +++ b/tests/integration/test_distributed_algorithm_it_contract.py @@ -3,6 +3,7 @@ from __future__ import annotations import ast +import json import shlex import subprocess import tomllib @@ -27,6 +28,7 @@ validate_output_dtype, validate_output_value, ) +from tools import algorithm_gate_provenance from tools.algorithm_gate_provenance import ( SourceRevision, build_preflight_provenance, @@ -35,6 +37,33 @@ write_provenance, ) + +def test_identity_manifest_export_is_an_exact_copy( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + source = ( + _ROOT / "tests" / "training" / "jobs" / "official_algorithm_identities.json" + ) + output = tmp_path / "identities.json" + monkeypatch.setattr( + "sys.argv", + [ + "algorithm_gate_provenance.py", + "export-identities", + "--manifest", + str(source), + "--output", + str(output), + "--expected-count", + "37", + ], + ) + + algorithm_gate_provenance.main() + + assert output.read_bytes() == source.read_bytes() + + _ROOT = Path(__file__).resolve().parents[2] _LOCAL = _ROOT / "scripts" / "run_distributed_algorithm_local_it.sh" _DISTRIBUTED = _ROOT / "scripts" / "run_distributed_algorithm_it.sh" @@ -354,6 +383,20 @@ def test_official_gate_matrix_covers_current_37_entry_points_once() -> None: for entry_point in entry_points } == set(ENTRY_POINT_DISTRIBUTIONS.items()) + manifest_path = ( + _ROOT / "tests" / "training" / "jobs" / "official_algorithm_identities.json" + ) + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + assert manifest["schema_version"] == 1 + assert manifest["entry_points"] == { + name: { + "algorithm_id": identity.algorithm_id, + "distribution": identity.distribution, + "implementation_id": identity.implementation_id, + } + for name, identity in sorted(OFFICIAL_ALGORITHM_IDENTITIES.items()) + } + gate = ( _ROOT / "tests" / "training" / "jobs" / "official_algorithm_gate_job.py" ).read_text(encoding="utf-8") diff --git a/tests/training/exporters/test_torch_recipe_source.py b/tests/training/exporters/test_torch_recipe_source.py index 8b4076c..59ebdf1 100644 --- a/tests/training/exporters/test_torch_recipe_source.py +++ b/tests/training/exporters/test_torch_recipe_source.py @@ -23,6 +23,7 @@ TorchStageRunIdentity, TorchStepResult, ) +from tributo.algorithms.api import AlgorithmConfigurationError from tributo.algorithms.spi import ( RayTorchAdapter, TorchArtifactContext, @@ -168,6 +169,9 @@ def _checkpoint(path: Path) -> Path: model = BinaryLinearRecipe().build_modules(None)["model"] assert isinstance(model, torch.nn.Module) torch.save(model.state_dict(), path / "model.pt") + (path / "metrics.json").write_text( + json.dumps({"train_loss": 0.5}), encoding="utf-8" + ) identity = TorchStageRunIdentity( "aabbccdd", "11223344", @@ -192,15 +196,13 @@ def _checkpoint(path: Path) -> Path: input_binding_digest="3" * 64, implementation_code_digest=CODE_DIGEST, payload_files={ - "model.pt": hashlib.sha256((path / "model.pt").read_bytes()).hexdigest() + name: hashlib.sha256((path / name).read_bytes()).hexdigest() + for name in ("model.pt", "metrics.json") }, ) (path / "torch_checkpoint_descriptor.json").write_text( json.dumps(descriptor.to_dict()), encoding="utf-8" ) - (path / "metrics.json").write_text( - json.dumps({"train_loss": 0.5}), encoding="utf-8" - ) return path @@ -234,7 +236,6 @@ def _adapter_checkpoint(path: Path) -> Path: }, adapter_identity=identity.implementation_id, resume_supported=False, - same_world_size_resume=None, ) (path / "torch_checkpoint_descriptor.json").write_text( json.dumps(descriptor.to_dict()), encoding="utf-8" @@ -275,6 +276,24 @@ def test_recipe_source_rejects_plan_identity_drift(tmp_path: Path) -> None: pass +def test_recipe_source_rejects_input_binding_digest_drift( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + provider = RayTorchSourceProvider() + options = _options().model_copy(update={"input_binding_digest": "4" * 64}) + monkeypatch.setattr( + "tributo.integrations.sources.ray_torch._load_reference", + lambda reference: (_ for _ in ()).throw( + AssertionError("implementation must not be loaded") + ), + ) + + with pytest.raises(AlgorithmConfigurationError, match="input binding digest"): + with provider.open_source(_checkpoint(tmp_path), options): + pass + + def test_recipe_source_uses_existing_onnx_bundle_pipeline(tmp_path: Path) -> None: checkpoint_dir = tmp_path / "checkpoint" checkpoint_dir.mkdir() diff --git a/tests/training/jobs/official_algorithm_gate_job.py b/tests/training/jobs/official_algorithm_gate_job.py index 15f4499..384a377 100644 --- a/tests/training/jobs/official_algorithm_gate_job.py +++ b/tests/training/jobs/official_algorithm_gate_job.py @@ -70,11 +70,6 @@ def _tensor_columns( f"model input {name!r} has unsupported dynamic trailing shape {shape}" ) width = math.prod(cast(tuple[int, ...], trailing)) if trailing else 1 - if len(trailing) > 1: - # Preserve higher-rank typed tensors as one vector-valued table column; - # flattening them into scalar columns would make inference reconstruct - # rank two and violate the manifest signature. - return (name,) return tuple(f"{name}__{index}" for index in range(width)) @@ -134,35 +129,15 @@ def _stage_bundle_inference_input( columns = _tensor_columns(name=field.name, shape=field.shape) projected.extend(columns) for column_index, column in enumerate(columns): - if len(field.shape[1:]) > 1: - trailing = tuple(cast(tuple[int, ...], field.shape[1:])) - width = math.prod(trailing) - values[column] = [ - np.asarray( - [ - _inference_value( - field.dtype, - row=row, - column=field_index + offset, - entry_point=entry_point, - ) - for offset in range(width) - ] - ) - .reshape(trailing) - .tolist() - for row in range(16) - ] - else: - values[column] = [ - _inference_value( - field.dtype, - row=row, - column=field_index + column_index, - entry_point=entry_point, - ) - for row in range(16) - ] + values[column] = [ + _inference_value( + field.dtype, + row=row, + column=field_index + column_index, + entry_point=entry_point, + ) + for row in range(16) + ] bindings.append( TensorInputBinding( tensor_name=field.name, @@ -1544,7 +1519,6 @@ def main() -> None: "ray": { "max_failures": 0, "storage_path": str(root / "autoencoder-ray-results"), - "resume": {"checkpoint_interval": 1}, }, "output": {"bundle_uri": str(root / "autoencoder-bundle")}, }, @@ -1572,7 +1546,6 @@ def main() -> None: "ray": { "max_failures": 0, "storage_path": str(root / "timeseries-ray-results"), - "resume": {"checkpoint_interval": 1}, }, "output": {"bundle_uri": str(root / "timeseries-bundle")}, }, @@ -1600,7 +1573,6 @@ def main() -> None: "ray": { "max_failures": 0, "storage_path": str(root / "lstm-ray-results"), - "resume": {"checkpoint_interval": 1}, }, "output": {"bundle_uri": str(root / "lstm-bundle")}, }, @@ -1628,7 +1600,6 @@ def main() -> None: "ray": { "max_failures": 0, "storage_path": str(root / "gru-ray-results"), - "resume": {"checkpoint_interval": 1}, }, "output": {"bundle_uri": str(root / "gru-bundle")}, }, @@ -1652,7 +1623,6 @@ def main() -> None: "ray": { "max_failures": 0, "storage_path": str(root / "dnn-ray-results"), - "resume": {"checkpoint_interval": 1}, }, "output": {"bundle_uri": str(root / "dnn-v2-bundle")}, }, @@ -1677,7 +1647,6 @@ def main() -> None: "ray": { "max_failures": 0, "storage_path": str(root / "pu-ray-results"), - "resume": {"checkpoint_interval": 1}, }, "output": {"bundle_uri": str(root / "pu-v2-bundle")}, }, @@ -1705,7 +1674,6 @@ def main() -> None: "ray": { "max_failures": 0, "storage_path": str(root / "two-tower-ray-results"), - "resume": {"checkpoint_interval": 1}, }, "output": {"bundle_uri": str(root / "two-tower-bundle")}, }, @@ -1760,7 +1728,6 @@ def main() -> None: "ray": { "max_failures": 0, "storage_path": str(root / "transformer-ray-results"), - "resume": {"checkpoint_interval": 1}, }, "output": {"bundle_uri": str(root / "transformer-bundle")}, }, diff --git a/tests/training/jobs/official_algorithm_identities.json b/tests/training/jobs/official_algorithm_identities.json new file mode 100644 index 0000000..94e49a0 --- /dev/null +++ b/tests/training/jobs/official_algorithm_identities.json @@ -0,0 +1,42 @@ +{ + "entry_points": { + "catboost.parallel_ensemble": {"algorithm_id": "catboost", "distribution": "tributo-algorithms-catboost", "implementation_id": "tributo.official.catboost.parallel_ensemble"}, + "difference_in_means_ate": {"algorithm_id": "difference_in_means_ate", "distribution": "tributo-algorithms-causal-core", "implementation_id": "tributo.official.causal.difference_in_means"}, + "dnn": {"algorithm_id": "dnn", "distribution": "tributo-algorithms-tabular-torch", "implementation_id": "tributo.official.tabular_torch.dnn"}, + "doubly_robust_ate": {"algorithm_id": "doubly_robust_ate", "distribution": "tributo-algorithms-causal-dr", "implementation_id": "tributo.official.causal_dr.aipw"}, + "dowhy_linear_refutation": {"algorithm_id": "dowhy_linear_refutation", "distribution": "tributo-algorithms-causal-dowhy", "implementation_id": "tributo.official.causal_dowhy.linear_refutation"}, + "extra_trees.joblib": {"algorithm_id": "extra_trees", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.extra_trees.joblib"}, + "extra_trees.native": {"algorithm_id": "extra_trees", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.extra_trees.native_ensemble"}, + "gcm_root_cause": {"algorithm_id": "gcm_root_cause", "distribution": "tributo-algorithms-causal-dowhy", "implementation_id": "tributo.official.causal_dowhy.gcm_root_cause"}, + "graphsage_node_classifier": {"algorithm_id": "graphsage_node_classifier", "distribution": "tributo-algorithms-graph-pyg", "implementation_id": "tributo.official.graph_pyg.graphsage"}, + "gru_classifier": {"algorithm_id": "gru_classifier", "distribution": "tributo-algorithms-timeseries", "implementation_id": "tributo.official.timeseries.gru"}, + "isolation_forest.parallel_ensemble": {"algorithm_id": "isolation_forest", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.classical.isolation_forest.parallel_ensemble"}, + "jagged_embedding_recommender": {"algorithm_id": "jagged_embedding_recommender", "distribution": "tributo-algorithms-recsys-torch", "implementation_id": "tributo.official.recsys_torch.jagged_embedding"}, + "kmeans.iterative": {"algorithm_id": "kmeans", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.classical.kmeans.iterative"}, + "kmeans_minibatch.iterative": {"algorithm_id": "kmeans_minibatch", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.classical.kmeans_minibatch.iterative"}, + "lightgbm.framework_native": {"algorithm_id": "lightgbm", "distribution": "tributo-algorithms-boosting", "implementation_id": "tributo.official.boosting.lightgbm"}, + "linear_dml_ate": {"algorithm_id": "linear_dml_ate", "distribution": "tributo-algorithms-causal-core", "implementation_id": "tributo.official.causal.linear_dml"}, + "linear_iv_ate": {"algorithm_id": "linear_iv_ate", "distribution": "tributo-algorithms-causal-core", "implementation_id": "tributo.official.causal.linear_iv"}, + "linear_regression.iterative": {"algorithm_id": "linear_regression", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.linear_regression.squared_l2"}, + "logistic_regression.iterative": {"algorithm_id": "logistic_regression", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.logistic_regression.binary_l2"}, + "lstm_classifier": {"algorithm_id": "lstm_classifier", "distribution": "tributo-algorithms-timeseries", "implementation_id": "tributo.official.timeseries.lstm"}, + "multinomial_nb": {"algorithm_id": "multinomial_nb", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.multinomial_nb.map_reduce"}, + "pc_stability_discovery": {"algorithm_id": "pc_stability_discovery", "distribution": "tributo-algorithms-causal-discovery", "implementation_id": "tributo.official.causal_discovery.pc_stability"}, + "pca.map_reduce": {"algorithm_id": "pca", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.classical.pca.map_reduce"}, + "pretrain_finetune_classifier": {"algorithm_id": "pretrain_finetune_classifier", "distribution": "tributo-algorithms-multistage-torch", "implementation_id": "tributo.official.multistage_torch.pretrain_finetune"}, + "pu": {"algorithm_id": "pu", "distribution": "tributo-algorithms-tabular-torch", "implementation_id": "tributo.official.tabular_torch.pu"}, + "random_forest.joblib": {"algorithm_id": "random_forest", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.random_forest.joblib"}, + "random_forest.native": {"algorithm_id": "random_forest", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.random_forest.native_ensemble"}, + "rgcn_node_classifier": {"algorithm_id": "rgcn_node_classifier", "distribution": "tributo-algorithms-graph-pyg", "implementation_id": "tributo.official.graph_pyg.rgcn"}, + "sgd_classifier.iterative": {"algorithm_id": "sgd_classifier", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.classical.sgd_classifier.iterative"}, + "sgd_regressor.iterative": {"algorithm_id": "sgd_regressor", "distribution": "tributo-algorithms-classical", "implementation_id": "tributo.official.classical.sgd_regressor.iterative"}, + "tabular_autoencoder": {"algorithm_id": "tabular_autoencoder", "distribution": "tributo-algorithms-representation", "implementation_id": "tributo.official.representation.tabular_autoencoder"}, + "teacher_student_distillation": {"algorithm_id": "teacher_student_distillation", "distribution": "tributo-algorithms-multistage-torch", "implementation_id": "tributo.official.multistage_torch.distillation"}, + "temporal_conv_classifier": {"algorithm_id": "temporal_conv_classifier", "distribution": "tributo-algorithms-timeseries", "implementation_id": "tributo.official.timeseries.temporal_conv"}, + "token_transformer_classifier": {"algorithm_id": "token_transformer_classifier", "distribution": "tributo-algorithms-transformers-nlp", "implementation_id": "tributo.official.transformer.token_classifier"}, + "two_tower_recommender": {"algorithm_id": "two_tower_recommender", "distribution": "tributo-algorithms-recsys-torch", "implementation_id": "tributo.official.recsys_torch.two_tower"}, + "x_learner.framework_native": {"algorithm_id": "x_learner", "distribution": "tributo-algorithms-causal-xlearner", "implementation_id": "tributo.official.causal_xlearner.xgboost"}, + "xgboost.framework_native": {"algorithm_id": "xgboost", "distribution": "tributo-algorithms-boosting", "implementation_id": "tributo.official.boosting.xgboost"} + }, + "schema_version": 1 +} diff --git a/tests/training/jobs/official_algorithm_matrix.py b/tests/training/jobs/official_algorithm_matrix.py index 2799c42..43ed0b6 100644 --- a/tests/training/jobs/official_algorithm_matrix.py +++ b/tests/training/jobs/official_algorithm_matrix.py @@ -2,8 +2,10 @@ from __future__ import annotations +import json from collections.abc import Mapping from dataclasses import dataclass +from pathlib import Path from types import MappingProxyType @@ -16,197 +18,46 @@ class OfficialAlgorithmIdentity: implementation_id: str -OFFICIAL_ALGORITHM_IDENTITIES: Mapping[str, OfficialAlgorithmIdentity] = ( - MappingProxyType( - { - "catboost.parallel_ensemble": OfficialAlgorithmIdentity( - "tributo-algorithms-catboost", - "catboost", - "tributo.official.catboost.parallel_ensemble", - ), - "difference_in_means_ate": OfficialAlgorithmIdentity( - "tributo-algorithms-causal-core", - "difference_in_means_ate", - "tributo.official.causal.difference_in_means", - ), - "dnn": OfficialAlgorithmIdentity( - "tributo-algorithms-tabular-torch", - "dnn", - "tributo.official.tabular_torch.dnn", - ), - "doubly_robust_ate": OfficialAlgorithmIdentity( - "tributo-algorithms-causal-dr", - "doubly_robust_ate", - "tributo.official.causal_dr.aipw", - ), - "dowhy_linear_refutation": OfficialAlgorithmIdentity( - "tributo-algorithms-causal-dowhy", - "dowhy_linear_refutation", - "tributo.official.causal_dowhy.linear_refutation", - ), - "extra_trees.joblib": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "extra_trees", - "tributo.official.extra_trees.joblib", - ), - "extra_trees.native": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "extra_trees", - "tributo.official.extra_trees.native_ensemble", - ), - "gcm_root_cause": OfficialAlgorithmIdentity( - "tributo-algorithms-causal-dowhy", - "gcm_root_cause", - "tributo.official.causal_dowhy.gcm_root_cause", - ), - "graphsage_node_classifier": OfficialAlgorithmIdentity( - "tributo-algorithms-graph-pyg", - "graphsage_node_classifier", - "tributo.official.graph_pyg.graphsage", - ), - "gru_classifier": OfficialAlgorithmIdentity( - "tributo-algorithms-timeseries", - "gru_classifier", - "tributo.official.timeseries.gru", - ), - "isolation_forest.parallel_ensemble": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "isolation_forest", - "tributo.official.classical.isolation_forest.parallel_ensemble", - ), - "jagged_embedding_recommender": OfficialAlgorithmIdentity( - "tributo-algorithms-recsys-torch", - "jagged_embedding_recommender", - "tributo.official.recsys_torch.jagged_embedding", - ), - "kmeans.iterative": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "kmeans", - "tributo.official.classical.kmeans.iterative", - ), - "kmeans_minibatch.iterative": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "kmeans_minibatch", - "tributo.official.classical.kmeans_minibatch.iterative", - ), - "lightgbm.framework_native": OfficialAlgorithmIdentity( - "tributo-algorithms-boosting", - "lightgbm", - "tributo.official.boosting.lightgbm", - ), - "linear_dml_ate": OfficialAlgorithmIdentity( - "tributo-algorithms-causal-core", - "linear_dml_ate", - "tributo.official.causal.linear_dml", - ), - "linear_iv_ate": OfficialAlgorithmIdentity( - "tributo-algorithms-causal-core", - "linear_iv_ate", - "tributo.official.causal.linear_iv", - ), - "linear_regression.iterative": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "linear_regression", - "tributo.official.linear_regression.squared_l2", - ), - "logistic_regression.iterative": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "logistic_regression", - "tributo.official.logistic_regression.binary_l2", - ), - "lstm_classifier": OfficialAlgorithmIdentity( - "tributo-algorithms-timeseries", - "lstm_classifier", - "tributo.official.timeseries.lstm", - ), - "multinomial_nb": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "multinomial_nb", - "tributo.official.multinomial_nb.map_reduce", - ), - "pc_stability_discovery": OfficialAlgorithmIdentity( - "tributo-algorithms-causal-discovery", - "pc_stability_discovery", - "tributo.official.causal_discovery.pc_stability", - ), - "pca.map_reduce": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "pca", - "tributo.official.classical.pca.map_reduce", - ), - "pretrain_finetune_classifier": OfficialAlgorithmIdentity( - "tributo-algorithms-multistage-torch", - "pretrain_finetune_classifier", - "tributo.official.multistage_torch.pretrain_finetune", - ), - "pu": OfficialAlgorithmIdentity( - "tributo-algorithms-tabular-torch", - "pu", - "tributo.official.tabular_torch.pu", - ), - "random_forest.joblib": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "random_forest", - "tributo.official.random_forest.joblib", - ), - "random_forest.native": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "random_forest", - "tributo.official.random_forest.native_ensemble", - ), - "rgcn_node_classifier": OfficialAlgorithmIdentity( - "tributo-algorithms-graph-pyg", - "rgcn_node_classifier", - "tributo.official.graph_pyg.rgcn", - ), - "sgd_classifier.iterative": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "sgd_classifier", - "tributo.official.classical.sgd_classifier.iterative", - ), - "sgd_regressor.iterative": OfficialAlgorithmIdentity( - "tributo-algorithms-classical", - "sgd_regressor", - "tributo.official.classical.sgd_regressor.iterative", - ), - "tabular_autoencoder": OfficialAlgorithmIdentity( - "tributo-algorithms-representation", - "tabular_autoencoder", - "tributo.official.representation.tabular_autoencoder", - ), - "teacher_student_distillation": OfficialAlgorithmIdentity( - "tributo-algorithms-multistage-torch", - "teacher_student_distillation", - "tributo.official.multistage_torch.distillation", - ), - "temporal_conv_classifier": OfficialAlgorithmIdentity( - "tributo-algorithms-timeseries", - "temporal_conv_classifier", - "tributo.official.timeseries.temporal_conv", - ), - "token_transformer_classifier": OfficialAlgorithmIdentity( - "tributo-algorithms-transformers-nlp", - "token_transformer_classifier", - "tributo.official.transformer.token_classifier", - ), - "two_tower_recommender": OfficialAlgorithmIdentity( - "tributo-algorithms-recsys-torch", - "two_tower_recommender", - "tributo.official.recsys_torch.two_tower", - ), - "x_learner.framework_native": OfficialAlgorithmIdentity( - "tributo-algorithms-causal-xlearner", - "x_learner", - "tributo.official.causal_xlearner.xgboost", - ), - "xgboost.framework_native": OfficialAlgorithmIdentity( - "tributo-algorithms-boosting", - "xgboost", - "tributo.official.boosting.xgboost", - ), - } - ) -) +def _load_official_algorithm_identities() -> Mapping[str, OfficialAlgorithmIdentity]: + path = Path(__file__).with_name("official_algorithm_identities.json") + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError( + "official algorithm identity manifest is unavailable" + ) from exc + entries = payload.get("entry_points") if isinstance(payload, Mapping) else None + if ( + not isinstance(payload, Mapping) + or payload.get("schema_version") != 1 + or not isinstance(entries, Mapping) + or len(entries) != 37 + ): + raise RuntimeError("official algorithm identity manifest is malformed") + identities: dict[str, OfficialAlgorithmIdentity] = {} + required = {"distribution", "algorithm_id", "implementation_id"} + for entry_point, value in entries.items(): + if ( + not isinstance(entry_point, str) + or not entry_point + or not isinstance(value, Mapping) + or set(value) != required + or any( + not isinstance(value[name], str) or not value[name] for name in required + ) + ): + raise RuntimeError( + "official algorithm identity manifest entry is malformed" + ) + identities[entry_point] = OfficialAlgorithmIdentity( + distribution=value["distribution"], + algorithm_id=value["algorithm_id"], + implementation_id=value["implementation_id"], + ) + return MappingProxyType(identities) + + +OFFICIAL_ALGORITHM_IDENTITIES = _load_official_algorithm_identities() def _build_distribution_entry_points() -> Mapping[str, tuple[str, ...]]: diff --git a/tests/training/test_portable_tune.py b/tests/training/test_portable_tune.py index 87a1a49..53cd67d 100644 --- a/tests/training/test_portable_tune.py +++ b/tests/training/test_portable_tune.py @@ -3,11 +3,13 @@ from __future__ import annotations from dataclasses import replace +from types import SimpleNamespace from tests.algorithms.conftest import map_reduce_registration, request_for from tributo.algorithms.api import ( AlgorithmOperation, AlgorithmRequest, + DistributionStrategy, ExecutionProfile, ExecutionRequest, ) @@ -18,6 +20,7 @@ from tributo.training.portable_tune import ( PortableTuneRunner, _fit_only_plan, + _trial_checkpoint_enabled, _trial_request, ) from tributo.training.tune_config import TuneSearchConfig @@ -77,6 +80,14 @@ def test_trial_request_applies_only_sampled_config_and_isolates_checkpoint() -> } +def test_torch_tune_trial_is_metric_only() -> None: + torch_plan = SimpleNamespace( + distribution_spec=SimpleNamespace(strategy=DistributionStrategy.RAY_TRAIN_TORCH) + ) + assert not _trial_checkpoint_enabled(torch_plan) + assert _trial_checkpoint_enabled(_resolved_plan()) + + def test_portable_tune_requires_tunable_descriptor_and_cluster_profile() -> None: registration = map_reduce_registration() descriptor = registration.spec diff --git a/tools/algorithm_gate_provenance.py b/tools/algorithm_gate_provenance.py index 34495a5..c77f25b 100644 --- a/tools/algorithm_gate_provenance.py +++ b/tools/algorithm_gate_provenance.py @@ -327,6 +327,10 @@ def _add_source_arguments(parser: argparse.ArgumentParser) -> None: def _parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser() commands = parser.add_subparsers(dest="command", required=True) + identities = commands.add_parser("export-identities") + identities.add_argument("--manifest", type=Path, required=True) + identities.add_argument("--output", type=Path, required=True) + identities.add_argument("--expected-count", type=int, required=True) preflight = commands.add_parser("preflight") _add_source_arguments(preflight) final = commands.add_parser("final") @@ -342,6 +346,32 @@ def _parser() -> argparse.ArgumentParser: def main() -> None: args = _parser().parse_args() + if args.command == "export-identities": + if args.output.exists(): + raise FileExistsError( + f"refusing to overwrite identity manifest: {args.output}" + ) + payload = json.loads(args.manifest.read_text(encoding="utf-8")) + entries = payload.get("entry_points") if isinstance(payload, Mapping) else None + if ( + not isinstance(payload, Mapping) + or payload.get("schema_version") != 1 + or not isinstance(entries, Mapping) + or len(entries) != args.expected_count + ): + raise ValueError("official algorithm identity manifest is malformed") + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_bytes(args.manifest.read_bytes()) + print( + json.dumps( + { + "entry_point_count": len(entries), + "sha256": _sha256_file(args.output), + }, + sort_keys=True, + ) + ) + return core = inspect_source(args.core_root) algorithms = inspect_source(args.algorithms_root) if args.command == "preflight":