spark_on_k8s.utils.driver_pod_template

Implementation reference. These symbols describe how the project works; their presence here does not make them supported application imports. Use the documented entry points and support policy.

load_driver_pod_template(template)

Load a local Pod YAML file, YAML text, or mapping without changing the input.

Source code in spark_on_k8s/utils/driver_pod_template.py

def load_driver_pod_template(template: str | Path | Mapping[str, Any]) -> dict[str, Any]:
    """Load a local Pod YAML file, YAML text, or mapping without changing the input."""
    from collections.abc import Mapping

    if isinstance(template, Mapping):
        pod = deepcopy(dict(template))
    else:
        if isinstance(template, Path):
            try:
                content = template.expanduser().read_text(encoding="utf-8")
            except OSError as exc:
                raise ValueError(f"Cannot read driver pod template file: {template}") from exc
        elif isinstance(template, str):
            try:
                path = Path(template).expanduser()
                is_file = "\n" not in template and path.is_file()
            except OSError:
                is_file = False
            if is_file:
                try:
                    content = path.read_text(encoding="utf-8")
                except OSError as exc:
                    raise ValueError(f"Cannot read driver pod template file: {template}") from exc
            else:
                content = template
        else:
            raise ValueError("Driver pod template must be YAML text, a local path, or a mapping")
        try:
            pod = yaml.load(content, Loader=_UniqueKeyLoader)
        except yaml.YAMLError as exc:
            raise ValueError("Driver pod template contains invalid YAML") from exc
    _validate_template(pod)
    return pod

merge_driver_pod_template(generated, template, driver_container_name=None)

Overlay a generated driver on a validated Pod template, preserving other fields.

The returned API dictionary preserves Kubernetes fields not modeled by this client version. Validation here covers template shape and merge ambiguity; the API server remains responsible for Kubernetes schema, admission, and cluster policy checks.

Source code in spark_on_k8s/utils/driver_pod_template.py

def merge_driver_pod_template(
    generated: k8s.V1PodTemplateSpec,
    template: dict[str, Any],
    driver_container_name: str | None = None,
) -> dict[str, Any]:
    """Overlay a generated driver on a validated Pod template, preserving other fields.

    The returned API dictionary preserves Kubernetes fields not modeled by this client
    version. Validation here covers template shape and merge ambiguity; the API server
    remains responsible for Kubernetes schema, admission, and cluster policy checks.
    """
    _validate_template(template)
    pod = deepcopy(template)
    with k8s.ApiClient() as api_client:
        managed = api_client.sanitize_for_serialization(generated)
    containers = pod["spec"]["containers"]
    names = [container["name"] for container in containers]
    selected = driver_container_name
    if selected is None:
        selected = "driver" if "driver" in names else names[0] if len(names) == 1 else None
    if selected not in names:
        raise ValueError(
            "Select an existing driver container with driver_container_name; "
            "selection is ambiguous or missing"
        )
    for container in containers:
        if container["name"] != selected and not container.get("image"):
            raise ValueError(f"Sidecar {container['name']!r} requires an image")
    driver = next(container for container in containers if container["name"] == selected)
    generated_driver = managed["spec"]["containers"][0]
    driver.pop("command", None)  # Spark images must use their driver entrypoint.
    for field in ("image", "imagePullPolicy", "args"):
        driver[field] = generated_driver[field]
    resources = driver.setdefault("resources", {})
    for field in ("requests", "limits"):
        resources[field] = {**resources.get(field, {}), **generated_driver["resources"].get(field, {})}
    for field, key in (("env", "name"), ("volumeMounts", "mountPath")):
        driver[field] = _merge_items(driver.get(field, []), generated_driver.get(field, []), _field_key(key))
    driver["ports"] = _merge_items(
        driver.get("ports", []),
        generated_driver.get("ports", []),
        lambda port: port.get("name") or (port.get("containerPort"), port.get("protocol", "TCP")),
    )
    driver["envFrom"] = deepcopy(driver.get("envFrom", []))
    for source in generated_driver.get("envFrom", []):
        if source not in driver["envFrom"]:
            driver["envFrom"].append(source)
    spec = pod["spec"]
    spec["containers"] = [driver, *(container for container in containers if container["name"] != selected)]
    spec.pop("serviceAccount", None)
    for field in ("serviceAccountName", "restartPolicy"):
        spec[field] = managed["spec"][field]
    # An explicit max_runtime wins; otherwise a template deadline is kept like other pod fields.
    if "activeDeadlineSeconds" in managed["spec"]:
        spec["activeDeadlineSeconds"] = managed["spec"]["activeDeadlineSeconds"]
    spec["nodeSelector"] = {**spec.get("nodeSelector", {}), **managed["spec"].get("nodeSelector", {})}
    for field, key in (("volumes", "name"), ("initContainers", "name"), ("hostAliases", "ip")):
        spec[field] = _merge_items(
            _items(spec.get(field, []), field), managed["spec"].get(field, []), _field_key(key)
        )
    spec["tolerations"] = _items(spec.get("tolerations", []), "tolerations")
    for toleration in managed["spec"].get("tolerations", []):
        if toleration not in spec["tolerations"]:
            spec["tolerations"].append(toleration)
    metadata = pod.setdefault("metadata", {})
    for field in ("name", "namespace"):
        metadata[field] = managed["metadata"][field]
    labels = metadata.setdefault("labels", {})
    for key in ("spark-app-name", "spark-app-id", "spark-role", "spark-ui-proxy"):
        labels.pop(key, None)
    labels.update(managed["metadata"]["labels"])
    metadata["annotations"] = {
        **metadata.get("annotations", {}),
        **managed["metadata"].get("annotations", {}),
        "kubectl.kubernetes.io/default-container": selected,
    }
    _validate_template(pod, allow_ownership=True)
    return pod