distributed-execution/src/distributed_execution/loosely_coupled/jobs.py

"""Loosely coupled reference implementations.

Dagster does not run the work here. It dispatches an external payload and then
acts as a listener: the payload writes structured messages over the
dagster-pipes protocol, which Dagster materialises as logs and metadata. The
payload never connects to the metadata database.

Two transports are provided. Both dispatch the *same* payload script and produce
the same output shape as the tightly coupled ``process_work_unit`` op, so the
graphs differ by exactly one node - see the guide's section 5.3.
"""

import os
import sys
from pathlib import Path

from dagster import (
    Failure,
    OpExecutionContext,
    Out,
    PipesSubprocessClient,
    graph,
    in_process_executor,
    multiprocess_executor,
    op,
)
from dagster_k8s import PipesK8sClient

from distributed_execution.ops import (
    generate_work_units,
    report_execution_target,
    summarise_results,
)

# Payload lives outside src/ so it can be built into its own image.
_REPO_ROOT = Path(__file__).resolve().parents[3]
PAYLOAD_SCRIPT = os.environ.get("PIPES_PAYLOAD_SCRIPT", str(_REPO_ROOT / "payload" / "work.py"))

PAYLOAD_IMAGE = os.environ.get(
    "PIPES_PAYLOAD_IMAGE",
    "code.europa.eu:4567/simpl/simpl-open/data/supporting-data-services"
    "/distributed-execution/distributed-execution/payload:0.0.1",
)
PAYLOAD_NAMESPACE = os.environ.get("PIPES_PAYLOAD_NAMESPACE", "dagster")

COMMON_TAGS = {
    "execution_target": "loosely_coupled",
    "business_operation": "DISTRIBUTED_EXECUTION_REFERENCE",
}


def _result_from_pipes(context: OpExecutionContext, completed) -> dict:
    """Read one unit's result off the message channel.

    A broken message path is the defining failure mode of this pattern: the
    external workload can exit 0 while reporting nothing, so silence is treated
    as a failure rather than an empty result.
    """
    messages = completed.get_custom_messages()
    if not messages:
        raise Failure(
            description=(
                "No pipes messages received from the external payload. The workload may have "
                "run successfully while its messages were lost. Verify the message path: for "
                "the Kubernetes transport, confirm the pod log stream is readable and not "
                "being truncated or intercepted by a log shipper."
            )
        )

    payload = messages[-1]
    leaked = payload.get("orchestration_env_visible") or []
    if leaked:
        context.log.warning(
            "Payload could see orchestration runtime credentials: %s. "
            "This defeats the isolation the loosely coupled target is chosen for.",
            ", ".join(leaked),
        )

    row = payload["results"][0]
    context.log.info("Unit %s computed by external worker %s", row["unit"], row["worker"])
    return row


@op(
    description="Dispatches one external payload as a local subprocess and listens on the pipes channel.",
    out=Out(dict),
)
def dispatch_external_work_subprocess(
    context: OpExecutionContext,
    unit: int,
    pipes_subprocess_client: PipesSubprocessClient,
) -> dict:
    completed = pipes_subprocess_client.run(
        context=context,
        command=[sys.executable, PAYLOAD_SCRIPT],
        extras={"units": [unit]},
    )
    return _result_from_pipes(context, completed)


@op(
    description="Dispatches one external payload as a Kubernetes Job and listens on the pod log stream.",
    out=Out(dict),
)
def dispatch_external_work_k8s(
    context: OpExecutionContext,
    unit: int,
    pipes_k8s_client: PipesK8sClient,
) -> dict:
    completed = pipes_k8s_client.run(
        context=context,
        image=PAYLOAD_IMAGE,
        command=["python", "/app/work.py"],
        namespace=PAYLOAD_NAMESPACE,
        extras={"units": [unit]},
        base_pod_meta={
            "labels": {
                "app.kubernetes.io/name": "distributed-execution-payload",
                "dagster/execution-target": "loosely-coupled",
            }
        },
    )
    return _result_from_pipes(context, completed)


@graph
def loosely_coupled_subprocess_reference():
    target_report = report_execution_target()
    units = generate_work_units()
    results = units.map(dispatch_external_work_subprocess).collect()
    return summarise_results(results, target_report)


@graph
def loosely_coupled_k8s_reference():
    target_report = report_execution_target()
    units = generate_work_units()
    results = units.map(dispatch_external_work_k8s).collect()
    return summarise_results(results, target_report)


loosely_coupled_subprocess_job = loosely_coupled_subprocess_reference.to_job(
    name="loosely_coupled_subprocess_job",
    description=(
        "Loosely coupled via subprocess transport. Demonstrates the pipes contract end to end "
        "on a laptop, with no cluster and no metadata database access from the payload."
    ),
    # Dispatch is I/O-bound and the real work happens in the payload, so orchestrating
    # in one process keeps the laptop demo free of platform-specific spawn behaviour.
    executor_def=in_process_executor,
    resource_defs={"pipes_subprocess_client": PipesSubprocessClient()},
    tags={**COMMON_TAGS, "transport": "subprocess"},
)

loosely_coupled_k8s_job = loosely_coupled_k8s_reference.to_job(
    name="loosely_coupled_k8s_job",
    description=(
        "Loosely coupled via Kubernetes Job dispatch. Messages return over the pod log stream. "
        "Requires a cluster and RBAC to create Jobs in the target namespace."
    ),
    executor_def=multiprocess_executor.configured({"max_concurrent": 2}),
    resource_defs={"pipes_k8s_client": PipesK8sClient()},
    tags={**COMMON_TAGS, "transport": "k8s_pod_logs"},
)

LOOSELY_COUPLED_JOBS = [
    loosely_coupled_subprocess_job,
    loosely_coupled_k8s_job,
]