pipeloom-databricks 0.5.0__py3-none-any.whl

pipeloom_databricks/grouping.py

@@ -0,0 +1,158 @@
+"""From a plan to the tasks of one pipeline's job: units, compute, groups and task keys."""
+
+from __future__ import annotations
+
+import posixpath
+from collections import defaultdict
+from dataclasses import dataclass
+
+from pipeloom.diagnostics import Report
+from pipeloom.graph.dag import find_cycles
+from pipeloom.graph.selectors import Selection, SelectorError, apply_to_plan
+from pipeloom.plan.model import Plan
+from pipeloom_databricks import keys
+from pipeloom_databricks.options import DatabricksDeployOptions
+
+
+@dataclass(frozen=True)
+class TaskUnit:
+    """One Databricks task: a unit, or a group of units."""
+
+    key: str
+    units: tuple[str, ...]
+    parents: tuple[str, ...]
+    """Keys of the tasks this one waits for, sorted."""
+    compute: str
+    """Key of an entry in the deploy block's `compute`."""
+
+
+def _select(plan: Plan, select: list[str], exclude: list[str]) -> frozenset[str]:
+    # A throwaway report: "matched nothing" is reported by the caller in its own terms.
+    return apply_to_plan(Selection(tuple(select), tuple(exclude)), plan, Report())
+
+
+def pipeline_units(
+    plan: Plan, opts: DatabricksDeployOptions, pipeline: str, report: Report
+) -> tuple[str, ...]:
+    """The units of a pipeline, sorted. Empty, with PLD004, for an unknown pipeline, a bad
+    selector or a selection that holds no unit."""
+    options = opts.pipelines.get(pipeline)
+    if options is None:
+        report.add(
+            "PLD004",
+            f"pipeline '{pipeline}' is not in the deploy block",
+            hint=f"pipelines: {', '.join(sorted(opts.pipelines))}",
+            context={"pipeline": pipeline},
+        )
+        return ()
+    try:
+        chosen = _select(plan, options.select, options.exclude)
+    except SelectorError as err:
+        report.add("PLD004", f"pipeline '{pipeline}': {err}", context={"pipeline": pipeline})
+        return ()
+    units = tuple(sorted(uid for uid, u in plan.units.items() if u.nodes[-1] in chosen))
+    if not units:
+        report.add(
+            "PLD004",
+            f"pipeline '{pipeline}' selects no unit",
+            context={"pipeline": pipeline},
+        )
+    return units
+
+
+def _rule_compute(plan: Plan, opts: DatabricksDeployOptions, report: Report) -> dict[str, str]:
+    """Persisted node id -> compute key, for the nodes a rule claims (the first rule wins)."""
+    claimed: dict[str, str] = {}
+    for i, rule in enumerate(opts.compute_rules):
+        try:
+            chosen = _select(plan, rule.select, rule.exclude)
+        except SelectorError as err:
+            report.add("PLD004", f"compute rule {i}: {err}", context={"rule": i})
+            continue
+        if not chosen:
+            report.add(
+                "PLD005",
+                f"compute rule {i} ({', '.join(rule.select)}) matches no model",
+                context={"rule": i},
+            )
+        for node in sorted(chosen):
+            claimed.setdefault(node, rule.compute)
+    return claimed
+
+
+def _group_of(plan: Plan, opts: DatabricksDeployOptions, unit: str) -> str:
+    if opts.task_grouping == "none":
+        return unit
+    node = plan.nodes[plan.units[unit].nodes[-1]]
+    if opts.task_grouping == "by_tag":
+        return f"tag_{node.tags[0] if node.tags else 'untagged'}"
+    return f"dir_{posixpath.dirname(node.file) or 'root'}"
+
+
+def build_tasks(
+    plan: Plan, opts: DatabricksDeployOptions, pipeline: str, report: Report
+) -> tuple[TaskUnit, ...]:
+    """The tasks of a pipeline's job, sorted by key, finalize not included.
+
+    Reports PLD001 (too many tasks), PLD003 (a cycle among groups), PLD004 (a group on two
+    computes) and PLD005. Returns what it could build; the caller checks `report`.
+    """
+    units = pipeline_units(plan, opts, pipeline, report)
+    claimed = _rule_compute(plan, opts, report)
+    members: dict[str, list[str]] = defaultdict(list)
+    for unit in units:
+        members[_group_of(plan, opts, unit)].append(unit)
+    group_of = {u: g for g, us in members.items() for u in us}
+
+    compute: dict[str, str] = {}
+    for group, us in sorted(members.items()):
+        used = sorted({claimed.get(plan.units[u].nodes[-1], opts.default_compute) for u in us})
+        if len(used) > 1:
+            report.add(
+                "PLD004",
+                f"task group '{group}' spans more than one compute: {', '.join(used)}",
+                hint="change task_grouping or the compute rules so a group has one compute",
+                context={"group": group, "compute": [c for c in used]},
+            )
+        compute[group] = used[0]
+
+    parents: dict[str, tuple[str, ...]] = {
+        g: tuple(
+            sorted({group_of[p] for u in us for p in plan.units[u].parents if p in group_of} - {g})
+        )
+        for g, us in members.items()
+    }
+    assigned = keys.assign_keys(sorted(members))
+    for cycle in find_cycles(parents):
+        path = [assigned[g] for g in cycle]
+        report.add(
+            "PLD003",
+            f"task groups form a cycle: {' -> '.join((*path, path[0]))}",
+            hint="group by another key, or set task_grouping to none",
+            context={"groups": [g for g in path]},
+        )
+
+    if len(members) + 1 > keys.MAX_TASKS:
+        report.add(
+            "PLD001",
+            f"pipeline '{pipeline}' needs {len(members) + 1} tasks with its finalize task, "
+            f"the limit is {keys.MAX_TASKS}",
+            hint="set task_grouping in the deploy block, or split the project into pipelines",
+            context={"pipeline": pipeline, "tasks": len(members) + 1},
+        )
+
+    tasks = [
+        TaskUnit(
+            assigned[g],
+            tuple(sorted(us)),
+            tuple(sorted(assigned[p] for p in parents[g])),
+            compute[g],
+        )
+        for g, us in members.items()
+    ]
+    return tuple(sorted(tasks, key=lambda t: t.key))
+
+
+def task_key_map(plan: Plan, opts: DatabricksDeployOptions, pipeline: str) -> dict[str, str]:
+    """Unit id -> the key of the task that runs it."""
+    return {u: t.key for t in build_tasks(plan, opts, pipeline, Report()) for u in t.units}

pipeloom_databricks/jobs_model.py

@@ -0,0 +1,116 @@
+"""The subset of the Jobs API that pipeloom renders, as frozen dataclasses.
+
+Field order is output order. `to_json` drops None and empty collections, and appends the keys of
+an `extra` mapping, sorted, after the declared fields.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field, fields, is_dataclass
+from typing import Literal
+
+from pydantic import JsonValue
+
+
+@dataclass(frozen=True)
+class JobParameter:
+    name: str
+    default: str
+
+
+@dataclass(frozen=True)
+class QueueSettings:
+    enabled: bool = True
+
+
+@dataclass(frozen=True)
+class EnvironmentSpec:
+    environment_version: str
+    dependencies: tuple[str, ...] = ()
+
+
+@dataclass(frozen=True)
+class JobEnvironment:
+    environment_key: str
+    spec: EnvironmentSpec
+
+
+@dataclass(frozen=True)
+class JobCluster:
+    job_cluster_key: str
+    new_cluster: Mapping[str, JsonValue]
+
+
+@dataclass(frozen=True)
+class PythonWheelTask:
+    package_name: str
+    entry_point: str
+    parameters: tuple[str, ...]
+
+
+@dataclass(frozen=True)
+class TaskDependency:
+    task_key: str
+
+
+@dataclass(frozen=True)
+class Task:
+    task_key: str
+    depends_on: tuple[TaskDependency, ...]
+    run_if: Literal["ALL_DONE"] | None
+    environment_key: str | None
+    job_cluster_key: str | None
+    libraries: tuple[Mapping[str, JsonValue], ...]
+    python_wheel_task: PythonWheelTask
+    max_retries: int | None
+    min_retry_interval_millis: int | None
+    timeout_seconds: int | None
+
+
+@dataclass(frozen=True)
+class JobSettings:
+    name: str
+    tags: Mapping[str, str]
+    parameters: tuple[JobParameter, ...]
+    max_concurrent_runs: int
+    queue: QueueSettings
+    environments: tuple[JobEnvironment, ...]
+    job_clusters: tuple[JobCluster, ...]
+    tasks: tuple[Task, ...]
+    extra: Mapping[str, JsonValue] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class SubmitSpec:
+    """The body of runs/submit: a one-time run, no job."""
+
+    run_name: str
+    tasks: tuple[Task, ...]
+    environments: tuple[JobEnvironment, ...]
+    queue: QueueSettings
+
+
+def _value(v: object) -> JsonValue:
+    if is_dataclass(v) and not isinstance(v, type):
+        return to_json(v)
+    if isinstance(v, (tuple, list)):
+        return [_value(x) for x in v]
+    if isinstance(v, Mapping):
+        return {str(k): _value(x) for k, x in v.items()}
+    return v  # type: ignore[return-value]  # str, int, float, bool, None
+
+
+def to_json(obj: object) -> dict[str, JsonValue]:
+    """The JSON object of a model dataclass, in field order."""
+    out: dict[str, JsonValue] = {}
+    extra: Mapping[str, JsonValue] = {}
+    for f in fields(obj):  # type: ignore[arg-type]
+        v = getattr(obj, f.name)
+        if f.name == "extra":
+            extra = v
+        elif v is not None and not (isinstance(v, (tuple, list, Mapping)) and not v):
+            out[f.name] = _value(v)
+    for key in sorted(extra):
+        out[key] = extra[key]
+    return out

pipeloom_databricks/keys.py

@@ -0,0 +1,45 @@
+"""Task keys: 1 to 100 characters of letters, digits, `_` and `-`."""
+
+from __future__ import annotations
+
+import hashlib
+import re
+from collections import Counter
+from collections.abc import Sequence
+from typing import Final
+
+FINALIZE_KEY: Final = "pipeloom__finalize"
+MAX_KEY_LEN: Final = 100
+MAX_TASKS: Final = 1000
+"""Read as `keys.MAX_TASKS` at call time, so a test can lower it."""
+
+_UNSAFE = re.compile(r"[^A-Za-z0-9_-]")
+_HASH_LEN = 8
+_KEEP = MAX_KEY_LEN - _HASH_LEN - 1
+
+
+def sanitize(raw: str) -> str:
+    """Replace every character outside `[A-Za-z0-9_-]` with `_`. Never empty."""
+    return _UNSAFE.sub("_", raw) or "_"
+
+
+def _hashed(key: str, raw: str) -> str:
+    digest = hashlib.sha256(raw.encode("utf-8")).hexdigest()[:_HASH_LEN]
+    return f"{key[:_KEEP]}_{digest}"
+
+
+def assign_keys(raw_ids: Sequence[str]) -> dict[str, str]:
+    """A task key for each raw id.
+
+    The key is the sanitized id. One longer than 100 characters, or one that several ids share
+    (or that is the finalize key), becomes its first 91 characters, `_` and 8 hex characters of
+    the sha256 of the raw id. Every id in a collision gets the suffix, so the result does not
+    depend on the order of `raw_ids`.
+    """
+    base = {raw: sanitize(raw) for raw in raw_ids}
+    taken = Counter(base.values())
+    taken[FINALIZE_KEY] += 1
+    return {
+        raw: _hashed(key, raw) if len(key) > MAX_KEY_LEN or taken[key] > 1 else key
+        for raw, key in base.items()
+    }

pipeloom_databricks/options.py

@@ -0,0 +1,66 @@
+"""The options of the `databricks` deploy block. No Databricks import."""
+
+from __future__ import annotations
+
+from typing import Annotated, Literal
+
+from pydantic import Field, JsonValue
+
+from pipeloom.config.types import Identifier
+from pipeloom.plugins import OptionsModel
+
+
+class ServerlessCompute(OptionsModel):
+    kind: Literal["serverless"] = "serverless"
+    environment_version: str = "5"
+    """Serverless environment version, as a string. Version 1 runs Python 3.10, below the floor."""
+    dependencies: list[str] | None = None
+    """pip requirement lines. Unset installs the pinned pipeloom distributions."""
+
+
+class JobClusterCompute(OptionsModel):
+    kind: Literal["job_cluster"]
+    new_cluster: dict[str, JsonValue]
+    """The `new_cluster` object of the Jobs API."""
+    libraries: list[dict[str, JsonValue]] | None = None
+    """Task libraries. Unset installs the pinned pipeloom distributions from PyPI."""
+
+
+Compute = Annotated[ServerlessCompute | JobClusterCompute, Field(discriminator="kind")]
+
+
+class ComputeRule(OptionsModel):
+    select: list[str]
+    """Selectors; the rule applies to the models they match."""
+    exclude: list[str] = Field(default_factory=list)
+    compute: Identifier
+    """Key of an entry in `compute`."""
+
+
+class PipelineOptions(OptionsModel):
+    select: list[str] = Field(default_factory=list)
+    """Selectors; empty means the whole plan."""
+    exclude: list[str] = Field(default_factory=list)
+
+
+def _main() -> dict[str, PipelineOptions]:
+    return {"main": PipelineOptions()}
+
+
+def _default() -> dict[str, Compute]:
+    return {"default": ServerlessCompute()}
+
+
+class DatabricksDeployOptions(OptionsModel):
+    pipelines: dict[Identifier, PipelineOptions] = Field(default_factory=_main)
+    """One job per pipeline."""
+    compute: dict[Identifier, Compute] = Field(default_factory=_default)
+    default_compute: Identifier = "default"
+    compute_rules: list[ComputeRule] = Field(default_factory=list)
+    """The first rule that matches a model wins."""
+    task_grouping: Literal["none", "by_tag", "by_folder"] = "none"
+    """none: one task per unit. by_tag: by the model's first tag. by_folder: by its folder."""
+    workspace_root: str | None = None
+    """Where a direct deploy uploads the project. Default /Workspace/Users/<me>/.pipeloom."""
+    settings: dict[str, JsonValue] = Field(default_factory=dict)
+    """Job settings passed through (schedule, notifications, run_as, tags ...)."""