sequence-base 0.3.0__py3-none-any.whl

sequence_base/schema_export.py

@@ -0,0 +1,238 @@
+"""
+JSON Schema export for every W0.1 wire type — the deliverable is "wire types + JSON schema".
+
+Each pydantic model is exported via `model_json_schema()`; `SeqFlags` (an IntFlag, for which
+`model_json_schema()` does not apply) is exported as an integer with a named-member enum. Every schema
+carries `x-contract-version` (constants.version.CONTRACT_VERSION) as a schema-level version stamp.
+
+Determinism is the contract: `render(...)` serializes with sorted keys + fixed indent, so the bytes are a
+pure function of the Python definitions. `export_all()` writes them under `sequence_base/schemas/`; a test
+regenerates in memory and asserts byte-for-byte equality with what is on disk (no drift), and that every
+instance serialized round-trips through its own exported schema.
+
+Regenerate after any shape change:  `python -m sequence_base.schema_export`
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from pydantic import BaseModel
+
+from sequence_base.constants.version import CONTRACT_VERSION
+from sequence_base.contract.accounts import (
+    BootstrapResponse,
+    SecretCreateRequest,
+    SecretInfo,
+    SecretRef,
+    UsageReport,
+)
+from sequence_base.contract.act import (
+    ActionChunk,
+    ActionStep,
+    ImageFrame,
+    Observation,
+    Proprioception,
+)
+from sequence_base.contract.benchmarks import (
+    BenchmarkRegisterRequest,
+    BenchmarkSpec,
+    Conditions,
+    RenderSpec,
+    Step,
+    Variant,
+)
+from sequence_base.contract.connect import ConnectRequest, ConnectResponse, Lease
+from sequence_base.contract.deployments import (
+    Autoscaler,
+    CodeBundle,
+    DeploymentAccepted,
+    DeploymentPolicy,
+    DeploymentRequest,
+    DeploymentStatusResponse,
+    Resources,
+    Timeouts,
+    VolumeAttach,
+)
+from sequence_base.contract.errors import ErrorBody, WarmingError
+from sequence_base.contract.eval import (
+    AdapterTransform,
+    ClampInfo,
+    Estimate,
+    EvalDryRunResponse,
+    EvalReport,
+    EvalRequest,
+    EvalResponse,
+    LossyTransform,
+    MetricAgg,
+    RolloutResult,
+    SuiteAgg,
+    TaskAgg,
+)
+from sequence_base.contract.flags import seqflags_json_schema
+from sequence_base.contract.policy import (
+    ActionContract,
+    Concurrency,
+    ImageSpec,
+    ImageStepApt,
+    ImageStepEnv,
+    ImageStepRun,
+    ImageStepUvPip,
+    PolicyRegistration,
+    PolicySpec,
+    VolumeMount,
+    VolumeRef,
+)
+
+SCHEMA_DIR = Path(__file__).parent / "schemas"
+
+# Every pydantic wire model, by the filename stem its schema is written under. One entry per shared
+# shape — this list IS the export manifest.
+WIRE_MODELS: dict[str, type[BaseModel]] = {
+    # contract 4 — /act data plane
+    "ImageFrame": ImageFrame,
+    "Proprioception": Proprioception,
+    "Observation": Observation,
+    "ActionStep": ActionStep,
+    "ActionChunk": ActionChunk,
+    # contract 1 — policy registry
+    "ImageStepUvPip": ImageStepUvPip,
+    "ImageStepApt": ImageStepApt,
+    "ImageStepEnv": ImageStepEnv,
+    "ImageStepRun": ImageStepRun,
+    "ImageSpec": ImageSpec,
+    "VolumeMount": VolumeMount,
+    "VolumeRef": VolumeRef,
+    "Concurrency": Concurrency,
+    "ActionContract": ActionContract,
+    "PolicySpec": PolicySpec,
+    "PolicyRegistration": PolicyRegistration,
+    # contract 2 — deployments
+    "Resources": Resources,
+    "Autoscaler": Autoscaler,
+    "Timeouts": Timeouts,
+    "VolumeAttach": VolumeAttach,
+    "CodeBundle": CodeBundle,
+    "DeploymentPolicy": DeploymentPolicy,
+    "DeploymentRequest": DeploymentRequest,
+    "DeploymentAccepted": DeploymentAccepted,
+    "DeploymentStatusResponse": DeploymentStatusResponse,
+    # contract 3 — connect / lease
+    "Lease": Lease,
+    "ConnectRequest": ConnectRequest,
+    "ConnectResponse": ConnectResponse,
+    # contract 5 — benchmark registration
+    "Variant": Variant,
+    "Conditions": Conditions,
+    "RenderSpec": RenderSpec,
+    "Step": Step,
+    "BenchmarkSpec": BenchmarkSpec,
+    "BenchmarkRegisterRequest": BenchmarkRegisterRequest,
+    # contract 5 — eval run / report / adapter
+    "EvalRequest": EvalRequest,
+    "ClampInfo": ClampInfo,
+    "Estimate": Estimate,
+    "EvalResponse": EvalResponse,
+    "EvalDryRunResponse": EvalDryRunResponse,
+    "RolloutResult": RolloutResult,
+    "SuiteAgg": SuiteAgg,
+    "TaskAgg": TaskAgg,
+    "MetricAgg": MetricAgg,
+    "LossyTransform": LossyTransform,
+    "EvalReport": EvalReport,
+    "AdapterTransform": AdapterTransform,
+    # contract 6 — accounts / secret / billing
+    "SecretRef": SecretRef,
+    "SecretCreateRequest": SecretCreateRequest,
+    "SecretInfo": SecretInfo,
+    "BootstrapResponse": BootstrapResponse,
+    "UsageReport": UsageReport,
+    # shared error bodies
+    "ErrorBody": ErrorBody,
+    "WarmingError": WarmingError,
+}
+
+
+def _stringify_int_property_names(schema: dict) -> None:
+    """JSON object keys are always strings, but pydantic renders an int-keyed dict (e.g.
+    dict[SeqFlags, str]) with integer `propertyNames` — inline or `$ref`-ed to an integer `$defs` enum —
+    which no serialized instance can satisfy (its keys are the strings "2", "16", …). Walk the schema and
+    rewrite any such `propertyNames` to the string form, so the exported schema accurately describes the
+    wire JSON. The referenced `$def` (the SeqFlags value form) is left untouched; only the key constraint
+    is inlined and string-encoded."""
+    defs = schema.get("$defs", {})
+
+    def resolve(pn: dict) -> dict | None:
+        if "$ref" in pn:
+            return defs.get(pn["$ref"].split("/")[-1])
+        return pn
+
+    def walk(node: object) -> None:
+        if isinstance(node, dict):
+            pn = node.get("propertyNames")
+            if isinstance(pn, dict):
+                target = resolve(pn)
+                if isinstance(target, dict) and target.get("type") == "integer":
+                    fixed: dict = {"type": "string"}
+                    if isinstance(target.get("enum"), list):
+                        fixed["enum"] = [str(v) for v in target["enum"]]
+                    node["propertyNames"] = fixed
+            for v in node.values():
+                walk(v)
+        elif isinstance(node, list):
+            for v in node:
+                walk(v)
+
+    walk(schema)
+
+
+def _reconcile_embedded_seqflags(schema: dict) -> None:
+    """`SeqFlags` is an IntFlag whose every member is its own distinct bit. pydantic's default
+    `model_json_schema()` emits an embedded `$defs.SeqFlags` (present whenever a model has a
+    dict[SeqFlags, ...] field, e.g. PolicyRegistration.bindings / BenchmarkSpec.bindings) as a plain
+    integer enum with NO `x-flag-names` — which diverges from the canonical standalone
+    `seqflags_json_schema()` (same values, plus the name→value map).
+
+    Route the embedded def through the single source: overwrite `$defs.SeqFlags` with the canonical
+    schema so it is byte-identical to standalone SeqFlags.json. This runs BEFORE
+    `_stringify_int_property_names`, so any `propertyNames: {$ref: #/$defs/SeqFlags}` that resolves against
+    this def also inherits the canonical value list — the key constraint stays in lockstep."""
+    defs = schema.get("$defs")
+    if isinstance(defs, dict) and "SeqFlags" in defs:
+        defs["SeqFlags"] = seqflags_json_schema(CONTRACT_VERSION)
+
+
+def build_schema(model: type[BaseModel]) -> dict:
+    """A model's JSON Schema, with the contract version stamped in."""
+    schema = model.model_json_schema()
+    _reconcile_embedded_seqflags(schema)
+    _stringify_int_property_names(schema)
+    schema["x-contract-version"] = CONTRACT_VERSION
+    return schema
+
+
+def all_schemas() -> dict[str, dict]:
+    """Every exported schema, by filename stem — pydantic models plus the SeqFlags IntFlag."""
+    out: dict[str, dict] = {name: build_schema(m) for name, m in WIRE_MODELS.items()}
+    out["SeqFlags"] = seqflags_json_schema(CONTRACT_VERSION)
+    return out
+
+
+def render(schema: dict) -> str:
+    """Deterministic bytes: sorted keys, fixed indent, trailing newline. A pure function of the schema."""
+    return json.dumps(schema, indent=2, sort_keys=True, ensure_ascii=False) + "\n"
+
+
+def export_all(out_dir: Path = SCHEMA_DIR) -> list[Path]:
+    out_dir.mkdir(parents=True, exist_ok=True)
+    written: list[Path] = []
+    for name, schema in all_schemas().items():
+        path = out_dir / f"{name}.json"
+        path.write_text(render(schema), encoding="utf-8")
+        written.append(path)
+    return written
+
+
+if __name__ == "__main__":
+    paths = export_all()
+    print(f"wrote {len(paths)} schemas to {SCHEMA_DIR} (contract_version={CONTRACT_VERSION})")

sequence_base/schemas/ActionChunk.json

@@ -0,0 +1,102 @@
+{
+  "$defs": {
+    "ActionSpace": {
+      "description": "Must be explicit: the same floats mean opposite things across spaces — executing an absolute\njoint position as a delta flings the arm into its limit stops. {schema.py:124}",
+      "enum": [
+        "joint_absolute",
+        "joint_delta",
+        "ee_absolute",
+        "ee_delta"
+      ],
+      "title": "ActionSpace",
+      "type": "string"
+    },
+    "ActionStep": {
+      "additionalProperties": false,
+      "properties": {
+        "index": {
+          "title": "Index",
+          "type": "integer"
+        },
+        "values": {
+          "items": {
+            "type": "number"
+          },
+          "title": "Values",
+          "type": "array"
+        }
+      },
+      "required": [
+        "index",
+        "values"
+      ],
+      "title": "ActionStep",
+      "type": "object"
+    }
+  },
+  "additionalProperties": false,
+  "description": "A chunk of actions. `covers_seconds` is what the robot schedules the next request against —\nit must ask again before this much motion finishes playing out.",
+  "properties": {
+    "action_dim": {
+      "title": "Action Dim",
+      "type": "integer"
+    },
+    "action_space": {
+      "$ref": "#/$defs/ActionSpace"
+    },
+    "control_frequency_hz": {
+      "title": "Control Frequency Hz",
+      "type": "number"
+    },
+    "covers_seconds": {
+      "title": "Covers Seconds",
+      "type": "number"
+    },
+    "open_loop_horizon": {
+      "anyOf": [
+        {
+          "type": "integer"
+        },
+        {
+          "type": "null"
+        }
+      ],
+      "default": null,
+      "title": "Open Loop Horizon"
+    },
+    "robot": {
+      "title": "Robot",
+      "type": "string"
+    },
+    "steps": {
+      "items": {
+        "$ref": "#/$defs/ActionStep"
+      },
+      "title": "Steps",
+      "type": "array"
+    },
+    "sustains_seconds": {
+      "anyOf": [
+        {
+          "type": "number"
+        },
+        {
+          "type": "null"
+        }
+      ],
+      "default": null,
+      "title": "Sustains Seconds"
+    }
+  },
+  "required": [
+    "steps",
+    "robot",
+    "action_space",
+    "action_dim",
+    "control_frequency_hz",
+    "covers_seconds"
+  ],
+  "title": "ActionChunk",
+  "type": "object",
+  "x-contract-version": "0.4.1"
+}

sequence_base/schemas/ActionContract.json

@@ -0,0 +1,43 @@
+{
+  "$defs": {
+    "ActionSpace": {
+      "description": "Must be explicit: the same floats mean opposite things across spaces — executing an absolute\njoint position as a delta flings the arm into its limit stops. {schema.py:124}",
+      "enum": [
+        "joint_absolute",
+        "joint_delta",
+        "ee_absolute",
+        "ee_delta"
+      ],
+      "title": "ActionSpace",
+      "type": "string"
+    }
+  },
+  "additionalProperties": false,
+  "description": "The fixed obs→ActionChunk contract a policy promises. Every field required; `action_space` reuses\ncontract 4's enum (not a free string). Each ActionChunk echoes exactly these invariants back.",
+  "properties": {
+    "action_dim": {
+      "title": "Action Dim",
+      "type": "integer"
+    },
+    "action_space": {
+      "$ref": "#/$defs/ActionSpace"
+    },
+    "control_frequency_hz": {
+      "title": "Control Frequency Hz",
+      "type": "number"
+    },
+    "robot": {
+      "title": "Robot",
+      "type": "string"
+    }
+  },
+  "required": [
+    "robot",
+    "action_space",
+    "action_dim",
+    "control_frequency_hz"
+  ],
+  "title": "ActionContract",
+  "type": "object",
+  "x-contract-version": "0.4.1"
+}

sequence_base/schemas/ActionStep.json

@@ -0,0 +1,23 @@
+{
+  "additionalProperties": false,
+  "properties": {
+    "index": {
+      "title": "Index",
+      "type": "integer"
+    },
+    "values": {
+      "items": {
+        "type": "number"
+      },
+      "title": "Values",
+      "type": "array"
+    }
+  },
+  "required": [
+    "index",
+    "values"
+  ],
+  "title": "ActionStep",
+  "type": "object",
+  "x-contract-version": "0.4.1"
+}

sequence_base/schemas/AdapterTransform.json

@@ -0,0 +1,42 @@
+{
+  "$defs": {
+    "AdapterKind": {
+      "enum": [
+        "lossless",
+        "lossy"
+      ],
+      "title": "AdapterKind",
+      "type": "string"
+    }
+  },
+  "additionalProperties": false,
+  "description": "One adapter transform, classified. A lossy transform must carry a non-empty reason; a lossless one\nneeds none. The lossy element's shape ({transform, reason}) aligns with EvalReport.adapter_lossy.",
+  "properties": {
+    "kind": {
+      "$ref": "#/$defs/AdapterKind"
+    },
+    "reason": {
+      "anyOf": [
+        {
+          "type": "string"
+        },
+        {
+          "type": "null"
+        }
+      ],
+      "default": null,
+      "title": "Reason"
+    },
+    "transform": {
+      "title": "Transform",
+      "type": "string"
+    }
+  },
+  "required": [
+    "kind",
+    "transform"
+  ],
+  "title": "AdapterTransform",
+  "type": "object",
+  "x-contract-version": "0.4.1"
+}

sequence_base/schemas/Autoscaler.json

@@ -0,0 +1,43 @@
+{
+  "additionalProperties": false,
+  "description": "= Modal AutoscalerSettings. Defaults mirror PolicySpec so an omitted block scale-to-zeroes.",
+  "properties": {
+    "buffer_containers": {
+      "anyOf": [
+        {
+          "type": "integer"
+        },
+        {
+          "type": "null"
+        }
+      ],
+      "default": null,
+      "title": "Buffer Containers"
+    },
+    "max_containers": {
+      "anyOf": [
+        {
+          "type": "integer"
+        },
+        {
+          "type": "null"
+        }
+      ],
+      "default": null,
+      "title": "Max Containers"
+    },
+    "min_containers": {
+      "default": 0,
+      "title": "Min Containers",
+      "type": "integer"
+    },
+    "scaledown_window": {
+      "default": 600,
+      "title": "Scaledown Window",
+      "type": "integer"
+    }
+  },
+  "title": "Autoscaler",
+  "type": "object",
+  "x-contract-version": "0.4.1"
+}