pipeloom-databricks 0.5.0__py3-none-any.whl

pipeloom_databricks/client/connect.py

@@ -0,0 +1,84 @@
+"""Choose the workspace connection (ADR-0036). No profile is ever picked automatically.
+
+Order: an explicit profile (`--profile`), a profile the user set in `DATABRICKS_CONFIG_PROFILE`,
+host and credentials in `DATABRICKS_HOST` and friends, runtime auth when
+`DATABRICKS_RUNTIME_VERSION` is set, else `ConnectionNotConfigured` (PLD010). Without a profile
+the SDK gets a `config_file` that does not exist, so it cannot fall back to a `DEFAULT` profile of
+`~/.databrickscfg`. The SDK loads lazily, only here and in `sdk.py`.
+"""
+
+from __future__ import annotations
+
+import importlib
+import tempfile
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any, Final
+
+from pipeloom.secrets.redaction import Redactor
+from pipeloom_databricks.client import ConnectionNotConfigured, DatabricksClient
+
+_CREDENTIALS: Final = {
+    "DATABRICKS_TOKEN": "token",
+    "DATABRICKS_CLIENT_ID": "client_id",
+    "DATABRICKS_CLIENT_SECRET": "client_secret",
+}
+_SECRET_ATTRS: Final = ("token", "client_secret", "password")
+
+
+def _factory(spec: str) -> DatabricksClient:
+    """`module:attr` called with no arguments (a test double; `--client-factory`)."""
+    module, sep, attr = spec.partition(":")
+    if not (module and sep and attr):
+        raise ConnectionNotConfigured(f"--client-factory '{spec}' is not MODULE:ATTR")
+    try:
+        make = getattr(importlib.import_module(module), attr)
+    except (ImportError, AttributeError) as err:
+        raise ConnectionNotConfigured(
+            f"--client-factory '{spec}' cannot be loaded: {err}"
+        ) from None
+    client: DatabricksClient = make()
+    return client
+
+
+def connect(
+    *, profile: str | None, env: Mapping[str, str], factory: str | None = None
+) -> DatabricksClient:
+    if factory:
+        return _factory(factory)
+    profile = profile or env.get("DATABRICKS_CONFIG_PROFILE") or None
+    host = env.get("DATABRICKS_HOST")
+    if profile:
+        kwargs: dict[str, Any] = {"profile": profile}
+    elif host:
+        kwargs = {"host": host}
+        kwargs.update({attr: env[name] for name, attr in _CREDENTIALS.items() if env.get(name)})
+    elif env.get("DATABRICKS_RUNTIME_VERSION"):
+        kwargs = {"auth_type": "runtime"}
+    else:
+        raise ConnectionNotConfigured(
+            "no Databricks connection: pass --profile, set DATABRICKS_CONFIG_PROFILE, set "
+            "DATABRICKS_HOST with credentials, or run on a Databricks runtime"
+        )
+
+    from databricks.sdk import WorkspaceClient
+    from databricks.sdk.core import Config
+
+    try:
+        if profile:
+            config = Config(**kwargs)
+        else:
+            with tempfile.TemporaryDirectory(prefix="pipeloom-") as folder:
+                config = Config(config_file=str(Path(folder) / "none.cfg"), **kwargs)
+        workspace = WorkspaceClient(config=config)
+    except ValueError as err:
+        raise ConnectionNotConfigured(f"the Databricks connection is invalid: {err}") from None
+
+    from pipeloom_databricks.client.sdk import SdkClient
+
+    redactor = Redactor()
+    for attr in _SECRET_ATTRS:
+        value = getattr(config, attr, None)
+        if isinstance(value, str):
+            redactor.register(value)
+    return SdkClient(workspace.api_client, str(config.host or "").rstrip("/"), redactor)

pipeloom_databricks/client/sdk.py

@@ -0,0 +1,275 @@
+"""`SdkClient`: the REST calls of the plugin, every one through `WorkspaceClient.api_client.do`.
+
+Going through `do` (not the typed SDK methods) keeps the JSON exactly as rendered, so a field the
+installed SDK does not know yet, such as `environment_version`, is never dropped. This is one of
+the three modules that import `databricks.sdk`; `connect` imports this one lazily.
+"""
+
+from __future__ import annotations
+
+import base64
+from collections.abc import Iterator, Mapping, Sequence
+from typing import Any, Final, Protocol
+from urllib.parse import quote
+
+from databricks.sdk.errors.base import DatabricksError
+from pydantic import JsonValue
+
+from pipeloom.secrets.redaction import Redactor
+from pipeloom_databricks.client import ApiError, DirEntry, JobSummary, RunInfo, TaskRunInfo
+
+JOBS: Final = "/api/2.2/jobs"
+WORKSPACE: Final = "/api/2.0/workspace"
+IMPORT_FORMAT: Final = "AUTO"
+"""The workspace import format; ADR-0036 explains why."""
+_RESULT: Final = {
+    "SUCCESS": "SUCCESS",
+    "SUCCESS_WITH_FAILURES": "SUCCESS_WITH_FAILURES",
+    "USER_CANCELED": "CANCELED",
+    "CANCELED": "CANCELED",
+    "UPSTREAM_FAILED": "UPSTREAM_FAILED",
+    "UPSTREAM_CANCELED": "UPSTREAM_CANCELED",
+    "MAX_CONCURRENT_RUNS_EXCEEDED": "MAXIMUM_CONCURRENT_RUNS_REACHED",
+    "EXCLUDED": "EXCLUDED",
+    "DISABLED": "DISABLED",
+}
+"""`termination_details.code` values that name a result state; any other code is a failure."""
+
+
+class _Api(Protocol):
+    def do(
+        self,
+        method: str,
+        path: str,
+        *,
+        query: dict[str, Any] | None = ...,
+        body: dict[str, Any] | None = ...,
+        raw: bool = ...,
+    ) -> Any: ...
+
+
+def _status(err: DatabricksError) -> int:
+    """The HTTP status of an SDK error: its own `status_code` when it has one, else the status
+    the SDK maps to the nearest class in its MRO (read from `STATUS_CODE_MAPPING`, so subclasses
+    such as `Aborted` or `DeadlineExceeded` resolve too), else 0."""
+    own = getattr(err, "status_code", None)
+    if isinstance(own, int) and not isinstance(own, bool):
+        return own
+    from databricks.sdk.errors import platform
+
+    by_class = {cls: code for code, cls in getattr(platform, "STATUS_CODE_MAPPING", {}).items()}
+    for cls in type(err).__mro__:
+        if cls in by_class:
+            return int(by_class[cls])
+    return 0
+
+
+def _states(item: Mapping[str, Any]) -> tuple[str, str | None, str | None, str | None]:
+    """life cycle state, result state, termination code and message of a run or a task run.
+
+    The `status` object is read when present, the legacy `state` object otherwise.
+    """
+    state: Mapping[str, Any] = item.get("state") or {}
+    status: Mapping[str, Any] = item.get("status") or {}
+    details: Mapping[str, Any] = status.get("termination_details") or {}
+    life = (status.get("state") if status else None) or state.get("life_cycle_state") or "UNKNOWN"
+    code = details.get("code")
+    result = state.get("result_state")
+    if result is None and life == "TERMINATED" and code:
+        result = _RESULT.get(code, "FAILED")
+    message = details.get("message") or state.get("state_message") or status.get("state_message")
+    return str(life), result, code, message or None
+
+
+def _task(raw: Mapping[str, Any]) -> TaskRunInfo:
+    life, result, code, message = _states(raw)
+    return TaskRunInfo(
+        task_key=str(raw["task_key"]),
+        run_id=int(raw["run_id"]),
+        attempt=int(raw.get("attempt_number") or 0),
+        life_cycle_state=life,
+        result_state=result,
+        termination_code=code,
+        message=message,
+    )
+
+
+def _run(raw: Mapping[str, Any], tasks: Sequence[Mapping[str, Any]] = ()) -> RunInfo:
+    life, result, _, _ = _states(raw)
+    params: dict[str, str] = {}
+    for p in raw.get("job_parameters") or []:
+        value = p.get("value", p.get("default"))
+        if value is not None:
+            params[str(p["name"])] = str(value)
+    repairs = [int(h["id"]) for h in raw.get("repair_history") or [] if h.get("id") is not None]
+    repair_states = {
+        int(h["id"]): _states(h)[0]
+        for h in raw.get("repair_history") or []
+        if h.get("id") is not None
+    }
+    job_id = raw.get("job_id")
+    start = raw.get("start_time")
+    return RunInfo(
+        run_id=int(raw["run_id"]),
+        job_id=None if job_id is None else int(job_id),
+        life_cycle_state=life,
+        result_state=result,
+        job_parameters=params,
+        tasks=tuple(_task(t) for t in tasks),
+        latest_repair_id=max(repairs) if repairs else None,
+        run_page_url=raw.get("run_page_url"),
+        start_time_ms=None if start is None else int(start),
+        repair_states=repair_states,
+    )
+
+
+class SdkClient:
+    """A `DatabricksClient` over the `api_client` of a `WorkspaceClient` (or a fake of it)."""
+
+    def __init__(self, api_client: _Api, host: str, redactor: Redactor | None = None) -> None:
+        self._api = api_client
+        self.host = host
+        self._redactor = redactor
+
+    def _call(
+        self,
+        method: str,
+        path: str,
+        *,
+        query: dict[str, Any] | None = None,
+        body: dict[str, Any] | None = None,
+        raw: bool = False,
+    ) -> Any:
+        try:
+            return self._api.do(method, path, query=query, body=body, raw=raw)
+        except DatabricksError as err:
+            text = str(err)
+            if self._redactor is not None:
+                text = self._redactor.mask(text)
+            raise ApiError(_status(err), text, err.error_code) from None
+
+    def list_jobs(self, *, name: str | None = None) -> Iterator[JobSummary]:
+        token: str | None = None
+        while True:
+            query: dict[str, Any] = {"limit": 100}
+            if name is not None:
+                query["name"] = name
+            if token:
+                query["page_token"] = token
+            page = self._call("GET", f"{JOBS}/list", query=query) or {}
+            for job in page.get("jobs") or []:
+                settings = job.get("settings") or {}
+                tags = {str(k): str(v) for k, v in (settings.get("tags") or {}).items()}
+                yield JobSummary(int(job["job_id"]), str(settings.get("name", "")), tags)
+            token = page.get("next_page_token")
+            if not token:
+                return
+
+    def create_job(self, settings: Mapping[str, JsonValue]) -> int:
+        out = self._call("POST", f"{JOBS}/create", body=dict(settings))
+        return int(out["job_id"])
+
+    def reset_job(self, job_id: int, settings: Mapping[str, JsonValue]) -> None:
+        self._call("POST", f"{JOBS}/reset", body={"job_id": job_id, "new_settings": dict(settings)})
+
+    def run_now(
+        self,
+        job_id: int,
+        *,
+        job_parameters: Mapping[str, str],
+        only: Sequence[str] | None,
+        idempotency_token: str | None,
+    ) -> int:
+        body: dict[str, Any] = {"job_id": job_id, "job_parameters": dict(job_parameters)}
+        if only is not None:
+            body["only"] = list(only)
+        if idempotency_token is not None:
+            body["idempotency_token"] = idempotency_token
+        return int(self._call("POST", f"{JOBS}/run-now", body=body)["run_id"])
+
+    def submit_run(self, spec: Mapping[str, JsonValue], *, idempotency_token: str | None) -> int:
+        body: dict[str, Any] = dict(spec)
+        if idempotency_token is not None:
+            body["idempotency_token"] = idempotency_token
+        return int(self._call("POST", f"{JOBS}/runs/submit", body=body)["run_id"])
+
+    def get_run(self, run_id: int) -> RunInfo:
+        query: dict[str, Any] = {"run_id": run_id, "include_history": "true"}
+        first = self._call("GET", f"{JOBS}/runs/get", query=query) or {}
+        tasks: list[Mapping[str, Any]] = list(first.get("tasks") or [])
+        token = first.get("next_page_token")
+        while token:
+            page = self._call("GET", f"{JOBS}/runs/get", query={**query, "page_token": token}) or {}
+            tasks.extend(page.get("tasks") or [])
+            token = page.get("next_page_token")
+        return _run(first, tasks)
+
+    def list_runs(self, job_id: int, *, active_only: bool = False) -> Iterator[RunInfo]:
+        token: str | None = None
+        while True:
+            query: dict[str, Any] = {"job_id": job_id, "limit": 25}
+            if active_only:
+                query["active_only"] = "true"
+            if token:
+                query["page_token"] = token
+            page = self._call("GET", f"{JOBS}/runs/list", query=query) or {}
+            for run in page.get("runs") or []:
+                yield _run(run, run.get("tasks") or [])
+            token = page.get("next_page_token")
+            if not token:
+                return
+
+    def repair_run(
+        self, run_id: int, *, rerun_tasks: Sequence[str], latest_repair_id: int | None
+    ) -> int:
+        body: dict[str, Any] = {
+            "run_id": run_id,
+            "rerun_tasks": list(rerun_tasks),
+            "rerun_dependent_tasks": True,
+        }
+        if latest_repair_id is not None:
+            body["latest_repair_id"] = latest_repair_id
+        return int(self._call("POST", f"{JOBS}/runs/repair", body=body)["repair_id"])
+
+    def current_user(self) -> str:
+        out = self._call("GET", "/api/2.0/preview/scim/v2/Me")
+        return str(out["userName"])
+
+    def mkdirs(self, path: str) -> None:
+        self._call("POST", f"{WORKSPACE}/mkdirs", body={"path": path})
+
+    def upload(self, path: str, data: bytes) -> None:
+        body = {
+            "path": path,
+            "format": IMPORT_FORMAT,
+            "content": base64.b64encode(data).decode("ascii"),
+            "overwrite": True,
+        }
+        self._call("POST", f"{WORKSPACE}/import", body=body)
+
+    def list_dir(self, path: str) -> list[DirEntry]:
+        out = self._call("GET", f"{WORKSPACE}/list", query={"path": path}) or {}
+        return [
+            DirEntry(
+                str(o["path"]), o.get("object_type") == "DIRECTORY", int(o.get("modified_at") or 0)
+            )
+            for o in out.get("objects") or []
+        ]
+
+    def delete(self, path: str, *, recursive: bool) -> None:
+        self._call("POST", f"{WORKSPACE}/delete", body={"path": path, "recursive": recursive})
+
+    def read_file(self, path: str) -> bytes:
+        out = self._call("GET", f"/api/2.0/fs/files{quote(path, safe='/')}", raw=True)
+        stream = out["contents"]
+        try:
+            data = stream.read()
+        finally:
+            close = getattr(stream, "close", None)
+            if close is not None:
+                close()
+        return bytes(data)
+
+    def get_secret(self, scope: str, key: str) -> str:
+        out = self._call("GET", "/api/2.0/secrets/get", query={"scope": scope, "key": key})
+        return str(out["value"])

pipeloom_databricks/deploy.py

@@ -0,0 +1,188 @@
+"""Direct deploy: upload the project, then create or reset the job of one pipeline.
+
+The job's parameters name the upload folder, which is named after the content hash, so a changed
+project changes the job settings and a second deploy of the same project is a no-op.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Mapping
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+from pipeloom.diagnostics import Report
+from pipeloom.plan.model import Plan
+from pipeloom.secrets.redaction import Redactor
+from pipeloom_databricks.bundle import Runtime
+from pipeloom_databricks.client import ApiError, DatabricksClient, JobSummary
+from pipeloom_databricks.jobs_model import to_json
+from pipeloom_databricks.render import Locations, RenderedJob, deploy_options, render_job
+from pipeloom_databricks.upload import MANIFEST_PATH, Payload, collect, prune, upload
+
+JOB_TAG = "pipeloom_job"
+HASH_TAG = "pipeloom_settings_hash"
+
+
+class DuplicateJobs(Exception):
+    """More than one job carries the same `pipeloom_job` tag (PLD006)."""
+
+    def __init__(self, job_tag: str, job_ids: tuple[int, ...]) -> None:
+        super().__init__(f"{len(job_ids)} jobs carry {JOB_TAG}={job_tag}: {list(job_ids)}")
+        self.job_tag = job_tag
+        self.job_ids = job_ids
+
+
+@dataclass(frozen=True)
+class DeployResult:
+    job_id: int
+    action: Literal["created", "updated", "unchanged"]
+    root: str
+    """The workspace folder that holds the uploaded project."""
+
+
+def find_job(
+    client: DatabricksClient, job_tag: str, *, name: str | None = None
+) -> JobSummary | None:
+    """The job tagged `pipeloom_job=<job_tag>`, or None.
+
+    With `name` the lookup by name goes first (cheap); when it finds no tagged job every job of
+    the workspace is scanned, because the tag, not the name, is the identity. Raises
+    `DuplicateJobs`.
+    """
+    for lookup in ([name] if name is not None else []) + [None]:
+        hits = [j for j in client.list_jobs(name=lookup) if j.tags.get(JOB_TAG) == job_tag]
+        if len(hits) > 1:
+            raise DuplicateJobs(job_tag, tuple(sorted(j.job_id for j in hits)))
+        if hits:
+            return hits[0]
+    return None
+
+
+@dataclass(frozen=True)
+class Prepared:
+    """Everything a deploy decides before it talks to the workspace."""
+
+    payload: Payload
+    parent: str
+    """`<workspace_root>/<project>/<pipeline>/<target>`: the folder of all uploads of the job."""
+    root: str
+    job: RenderedJob
+    locations: Locations
+
+
+def prepare(
+    project_dir: Path,
+    plan: Plan,
+    pipeline: str,
+    runtime: Runtime,
+    report: Report,
+    *,
+    user: Callable[[], str],
+    versions: Mapping[str, str] | None = None,
+) -> Prepared | None:
+    """Collect the files and render the job for their upload folder. Touches no workspace except
+    through `user`, which is called only when `workspace_root` is not set."""
+    opts = deploy_options(plan, report)
+    if opts is None:
+        return None
+    payload = collect(project_dir, plan, runtime.artifacts_root, report)
+    if payload is None:
+        return None
+    meta = plan.metadata
+    base = opts.workspace_root or f"/Workspace/Users/{user()}/.pipeloom"
+    parent = f"{base.rstrip('/')}/{meta.project}/{pipeline}/{meta.target}"
+    root = f"{parent}/{payload.digest[:12]}"
+    locs = Locations(
+        project_root=root,
+        manifest_path=f"{root}/{MANIFEST_PATH}",
+        artifacts_root=runtime.artifacts_root,
+        state_shared=runtime.state_shared,
+        cancel_grace=runtime.cancel_grace,
+    )
+    job = render_job(plan, pipeline, locs, report, versions=versions)
+    return None if job is None else Prepared(payload, parent, root, job, locs)
+
+
+@contextmanager
+def api_errors(report: Report, redactor: Redactor | None = None) -> Iterator[None]:
+    """Turn `DuplicateJobs` into PLD006 and `ApiError` into PLD009 in `report`; the block then
+    ends early (the caller returns None)."""
+    try:
+        yield
+    except DuplicateJobs as err:
+        report.add(
+            "PLD006",
+            f"{len(err.job_ids)} jobs carry {JOB_TAG}={err.job_tag}: ids {list(err.job_ids)}",
+            hint="delete or untag the extra jobs",
+        )
+    except ApiError as err:
+        text = redactor.mask(err.message) if redactor else err.message
+        code = f", {err.error_code}" if err.error_code else ""
+        report.add(
+            "PLD009",
+            f"Databricks API request failed (HTTP {err.status}{code}): {text}",
+            context={"status": err.status},
+        )
+
+
+def deploy_prepared(
+    client: DatabricksClient,
+    prepared: Prepared,
+    existing: JobSummary | None,
+    *,
+    report: Report,
+    redactor: Redactor | None = None,
+) -> DeployResult:
+    """Upload and create or reset the job of `prepared`; `existing` is the job found by tag.
+
+    Raises `ApiError`; wrap the call in `api_errors`.
+    """
+    job, root = prepared.job, prepared.root
+    if existing is not None and existing.tags.get(HASH_TAG) == job.settings_hash:
+        return DeployResult(existing.job_id, "unchanged", root)
+    upload(client, prepared.payload, root)
+    settings = to_json(job.settings)
+    if existing is None:
+        result = DeployResult(client.create_job(settings), "created", root)
+    else:
+        client.reset_job(existing.job_id, settings)
+        result = DeployResult(existing.job_id, "updated", root)
+    prune(
+        client,
+        prepared.parent,
+        prepared.payload.digest[:12],
+        report,
+        mask=redactor.mask if redactor else str,
+    )
+    return result
+
+
+def deploy(
+    client: DatabricksClient,
+    project_dir: Path,
+    plan: Plan,
+    pipeline: str,
+    runtime: Runtime,
+    *,
+    report: Report,
+    redactor: Redactor | None = None,
+    versions: Mapping[str, str] | None = None,
+) -> DeployResult | None:
+    """Deploy one pipeline, or return None with PLD004, PLD006, PLD009 or PLD012 in `report`."""
+    with api_errors(report, redactor):
+        prepared = prepare(
+            project_dir,
+            plan,
+            pipeline,
+            runtime,
+            report,
+            user=client.current_user,
+            versions=versions,
+        )
+        if prepared is None:
+            return None
+        existing = find_job(client, prepared.job.job_tag, name=prepared.job.settings.name)
+        return deploy_prepared(client, prepared, existing, report=report, redactor=redactor)
+    return None

pipeloom_databricks/finalize_states.py

@@ -0,0 +1,72 @@
+"""What the Jobs API knows about a run's tasks, as the unit states core finalize accepts.
+
+A unit whose task died (a lost cluster, a timeout, a cancel) wrote no attempt file. The finalize
+task asks the Jobs API how each task ended and hands the answer to `pipeloom finalize
+--unit-states`, so the run fails (or is cancelled) instead of looking like a success.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Collection, Mapping
+
+from pipeloom.runtime.finalize_run import UNIT_STATES_SCHEMA, ExternalUnitState
+from pipeloom_databricks.client import DatabricksClient, RunInfo, TaskRunInfo
+
+
+def task_states(client: DatabricksClient, job_run_id: int) -> Mapping[str, TaskRunInfo]:
+    """The tasks of a job run by task key (the latest attempt of each)."""
+    return {t.task_key: t for t in client.get_run(job_run_id).tasks}
+
+
+def unit_states(
+    key_map: Mapping[str, str], run: RunInfo, with_attempts: Collection[str]
+) -> dict[str, ExternalUnitState]:
+    """The external state of every unit that has no attempt file.
+
+    `key_map` is unit id -> task key. EXCLUDED is `excluded`; UPSTREAM_FAILED and
+    UPSTREAM_CANCELED are `skipped` (upstream_failed); CANCELED and TIMEDOUT are `cancelled`; FAILED
+    is `failed` with PLD011 and the termination message. Any other state (success, still running,
+    unknown task) gives no entry, so core finalize counts the unit as not run.
+    """
+    tasks = {t.task_key: t for t in run.tasks}
+    out: dict[str, ExternalUnitState] = {}
+    for unit, key in sorted(key_map.items()):
+        task = tasks.get(key)
+        if unit in with_attempts or task is None:
+            continue
+        ids = {
+            "databricks_task_key": task.task_key,
+            "databricks_task_run_id": str(task.run_id),
+        }
+        match task.result_state:
+            case "EXCLUDED":
+                out[unit] = ExternalUnitState("excluded")
+            case "UPSTREAM_FAILED" | "UPSTREAM_CANCELED":
+                out[unit] = ExternalUnitState("skipped", orchestrator_ids=ids)
+            case "CANCELED" | "TIMEDOUT":
+                out[unit] = ExternalUnitState("cancelled", orchestrator_ids=ids)
+            case "FAILED":
+                reason = task.message or task.termination_code or "no reason given"
+                out[unit] = ExternalUnitState(
+                    "failed",
+                    code="PLD011",
+                    message=f"task {task.task_key} ended without a unit result: {reason}",
+                    orchestrator_ids=ids,
+                )
+    return out
+
+
+def unit_states_json(states: Mapping[str, ExternalUnitState]) -> bytes:
+    """The `pipeloom.unit_states/1` file for `states`."""
+    units: dict[str, dict[str, object]] = {}
+    for unit, state in sorted(states.items()):
+        doc: dict[str, object] = {"status": state.status}
+        for name in ("reason", "code", "message"):
+            if getattr(state, name) is not None:
+                doc[name] = getattr(state, name)
+        if state.orchestrator_ids:
+            doc["orchestrator_ids"] = dict(state.orchestrator_ids)
+        units[unit] = doc
+    doc_text = json.dumps({"schema": UNIT_STATES_SCHEMA, "units": units}, ensure_ascii=True)
+    return (doc_text + "\n").encode("ascii")