(wip) flyte backfill <trigger> for scheduled triggers

src/flyte/backfill/__init__.py

@@ -0,0 +1,33 @@
+"""Backfill a scheduled trigger.
+
+Re-runs the slots a cron or fixed-rate trigger would have fired over a window,
+naming each run exactly the way a real fire would so that slots which already ran
+are recognised and skipped -- unless the backfill is forced, which salts the names
+into a separate namespace and re-runs them.
+
+The runs are created by a driver run inside the cluster, not from the client. See
+``flyte backfill --help``.
+"""
+
+from ._driver import backfill_driver, launch_backfill
+from ._execute import execute_plan, probe_existing
+from ._naming import candidate_run_names, scheduled_run_name
+from ._plan import DEFAULT_MAX_RUNS, BackfillPlan, BackfillSlot, build_plan, occurrences_for
+from ._schedule import CronParseError, cron_occurrences, fixed_rate_occurrences
+
+__all__ = [
+    "DEFAULT_MAX_RUNS",
+    "BackfillPlan",
+    "BackfillSlot",
+    "CronParseError",
+    "backfill_driver",
+    "build_plan",
+    "candidate_run_names",
+    "cron_occurrences",
+    "execute_plan",
+    "fixed_rate_occurrences",
+    "launch_backfill",
+    "occurrences_for",
+    "probe_existing",
+    "scheduled_run_name",
+]

src/flyte/backfill/_driver.py

@@ -0,0 +1,152 @@
+"""The run that performs a backfill.
+
+``flyte backfill`` does not create the backfilled runs from your machine. It
+launches one small driver run that creates them from inside the cluster, so a
+long backfill does not depend on a laptop staying connected, and so the work is
+itself observable, retryable, and abortable like any other run.
+
+Every run the driver creates is linked back to it as a spawned child, so the
+backfill shows up as one tree rather than a scatter of unrelated runs.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from flyte.remote import Run
+
+    from ._plan import BackfillPlan
+
+__all__ = ["backfill_driver", "encode_plan", "launch_backfill"]
+
+
+@dataclass
+class _EncodedPlan:
+    """The approved plan, in the form the driver receives it.
+
+    The exact slot list is carried across rather than recomputed, so the driver
+    creates precisely what was shown and approved -- no re-expansion, no drift if
+    the trigger's schedule changes in between.
+    """
+
+    org: str
+    project: str
+    domain: str
+    task_name: str
+    trigger_name: str
+    force: bool
+    salt: str | None
+    slots: list[dict]
+
+
+def encode_plan(plan: "BackfillPlan") -> str:
+    payload = _EncodedPlan(
+        org=plan.org,
+        project=plan.project,
+        domain=plan.domain,
+        task_name=plan.task_name,
+        trigger_name=plan.trigger_name,
+        force=plan.force,
+        salt=plan.salt,
+        slots=[
+            {
+                "at": s.scheduled_at.isoformat(),
+                "name": s.run_name,
+                "already_ran": s.already_ran,
+            }
+            for s in plan.to_create
+        ],
+    )
+    return json.dumps(payload.__dict__)
+
+
+def _decode_plan(encoded: str) -> tuple["BackfillPlan", list]:
+    from ._plan import BackfillPlan, BackfillSlot
+
+    raw = json.loads(encoded)
+    slots = [
+        BackfillSlot(
+            scheduled_at=datetime.fromisoformat(s["at"]),
+            run_name=s["name"],
+            already_ran=s["already_ran"],
+        )
+        for s in raw["slots"]
+    ]
+    plan = BackfillPlan(
+        trigger_name=raw["trigger_name"],
+        task_name=raw["task_name"],
+        project=raw["project"],
+        domain=raw["domain"],
+        org=raw["org"],
+        schedule="",
+        start=slots[0].scheduled_at if slots else datetime.now(timezone.utc),
+        end=slots[-1].scheduled_at if slots else datetime.now(timezone.utc),
+        force=raw["force"],
+        salt=raw["salt"],
+        queue=None,
+        max_runs=len(slots),
+        slots=slots,
+    )
+    return plan, slots
+
+
+async def backfill_driver(encoded_plan: str) -> str:
+    """Create every run in the encoded plan. Runs inside the cluster.
+
+    Returns a short summary. Individual slot failures are reported rather than
+    raised, so one rejected slot does not abandon the rest of the backfill.
+    """
+    from flyte.remote import Trigger
+
+    from ._execute import execute_plan
+
+    plan, _ = _decode_plan(encoded_plan)
+    details = await Trigger.get.aio(name=plan.trigger_name, task_name=plan.task_name)
+
+    created = 0
+    deduped = 0
+    failed: list[str] = []
+
+    def _record(result) -> None:
+        nonlocal created, deduped
+        if result.error:
+            failed.append(f"{result.slot.scheduled_at.isoformat()}: {result.error}")
+        elif result.created:
+            created += 1
+        else:
+            deduped += 1
+
+    await execute_plan(plan, details, on_slot=_record)
+
+    summary = f"backfill {plan.trigger_name}: {created} created, {deduped} already existed, {len(failed)} failed"
+    if failed:
+        summary += "\n" + "\n".join(failed[:20])
+    print(summary, flush=True)
+    return summary
+
+
+def launch_backfill(plan: "BackfillPlan", name: str | None = None) -> "Run":
+    """Launch the driver run that performs ``plan``."""
+    import flyte
+
+    env = flyte.TaskEnvironment(
+        name="flyte-backfill",
+        resources=flyte.Resources(cpu=1, memory="500Mi"),
+    )
+    driver = env.task(backfill_driver)
+    run = flyte.with_runcontext(
+        mode="remote",
+        name=name,
+        # The driver has no local source file of its own; ship it as a code bundle.
+        interactive_mode=True,
+        queue=plan.queue,
+    ).run(driver, encode_plan(plan))
+    from typing import cast
+
+    from flyte.remote import Run as _Run
+
+    return cast(_Run, run)

src/flyte/backfill/_execute.py

@@ -0,0 +1,170 @@
+"""Create the runs a backfill plan describes.
+
+A backfilled slot is created the same way the scheduler creates a real fire: the
+run is addressed by ``TriggerName`` (so the control plane resolves the task,
+inputs and run spec from the trigger itself), named by the deterministic slot
+hash, and marked as schedule-sourced.
+
+Using the schedule-trigger source matters. The control plane rewrites the run
+name's prefix for automation-sourced runs, and de-duplication happens against the
+rewritten name. Creating these runs under any other source would store them under
+a different name than a real fire, and the de-duplication that makes backfill safe
+to re-run would silently stop working.
+"""
+
+from __future__ import annotations
+
+from datetime import timezone
+from typing import TYPE_CHECKING, Callable
+
+if TYPE_CHECKING:
+    from ._plan import BackfillPlan, BackfillSlot
+
+__all__ = ["SlotResult", "execute_plan"]
+
+
+class SlotResult:
+    """Outcome of creating one slot's run."""
+
+    __slots__ = ("created", "error", "run_name", "slot")
+
+    def __init__(self, slot: "BackfillSlot", run_name: str, created: bool, error: str | None = None):
+        self.slot = slot
+        self.run_name = run_name
+        #: False when the control plane returned a pre-existing run instead of creating one.
+        self.created = created
+        self.error = error
+
+
+def _kickoff_arg_name(details) -> str:
+    """The input variable a trigger binds its scheduled time to, if any."""
+    schedule = details.pb2.automation_spec.schedule
+    return schedule.kickoff_time_input_arg or ""
+
+
+def _build_request(
+    plan: "BackfillPlan",
+    slot: "BackfillSlot",
+    kickoff_arg: str,
+):
+    from flyteidl2.common import identifier_pb2
+    from flyteidl2.core import literals_pb2
+    from flyteidl2.task import common_pb2 as task_common_pb2
+    from flyteidl2.workflow import run_definition_pb2, run_service_pb2
+    from google.protobuf import timestamp_pb2
+
+    scheduled = slot.scheduled_at
+    if scheduled.tzinfo is None:
+        scheduled = scheduled.replace(tzinfo=timezone.utc)
+    ts = timestamp_pb2.Timestamp()
+    ts.FromDatetime(scheduled.astimezone(timezone.utc))
+
+    req = run_service_pb2.CreateRunRequest(
+        run_id=identifier_pb2.RunIdentifier(
+            org=plan.org,
+            project=plan.project,
+            domain=plan.domain,
+            name=slot.run_name,
+        ),
+        trigger_name=identifier_pb2.TriggerName(
+            org=plan.org,
+            project=plan.project,
+            domain=plan.domain,
+            task_name=plan.task_name,
+            name=plan.trigger_name,
+        ),
+        source=run_definition_pb2.RUN_SOURCE_SCHEDULE_TRIGGER,
+        run_start_time=ts,
+    )
+
+    # Only triggers that bind their scheduled time to an input need the literal;
+    # everything else reads the time from run_start_time.
+    if kickoff_arg:
+        req.inputs.CopyFrom(
+            task_common_pb2.Inputs(
+                literals=[
+                    task_common_pb2.NamedLiteral(
+                        name=kickoff_arg,
+                        value=literals_pb2.Literal(
+                            scalar=literals_pb2.Scalar(
+                                primitive=literals_pb2.Primitive(datetime=ts),
+                            )
+                        ),
+                    )
+                ]
+            )
+        )
+    return req
+
+
+async def execute_plan(
+    plan: "BackfillPlan",
+    details,
+    on_slot: Callable[[SlotResult], None] | None = None,
+) -> list[SlotResult]:
+    """Create a run for every slot in ``plan.to_create``.
+
+    Slots the plan marked as already run are skipped unless the plan is forced.
+    Creating a run whose name already exists is not an error -- the control plane
+    returns the existing run -- so a re-run of the same backfill is a no-op.
+    """
+    from flyte._initialize import ensure_client, get_client
+
+    ensure_client()
+    kickoff_arg = _kickoff_arg_name(details)
+    results: list[SlotResult] = []
+
+    for slot in plan.to_create:
+        req = _build_request(plan, slot, kickoff_arg)
+        try:
+            resp = await get_client().run_service.create_run(req)
+            returned = resp.run.action.id.run.name or slot.run_name
+            # The server hands back the pre-existing run when the name is taken,
+            # which is the de-duplication path rather than a failure.
+            created = returned == slot.run_name or not slot.already_ran
+            result = SlotResult(slot, returned, created)
+        except Exception as exc:  # surfaced per slot; one bad slot must not sink the rest
+            result = SlotResult(slot, slot.run_name, False, error=str(exc))
+        results.append(result)
+        if on_slot is not None:
+            on_slot(result)
+    return results
+
+
+async def probe_existing(
+    plan_names: list[str],
+) -> set[str]:
+    """Return which of ``plan_names`` already exist as runs.
+
+    Used to mark slots as already run before anything is created, so the preview
+    and the confirmation prompt reflect what will actually happen.
+    """
+    from flyteidl2.workflow import run_service_pb2
+
+    from flyte._initialize import ensure_client, get_client
+
+    ensure_client()
+    from flyte._initialize import get_init_config
+
+    cfg = get_init_config()
+    found: set[str] = set()
+    for name in plan_names:
+        try:
+            from flyteidl2.common import identifier_pb2
+
+            await get_client().run_service.get_run_details(
+                run_service_pb2.GetRunDetailsRequest(
+                    run_id=identifier_pb2.RunIdentifier(
+                        org=cfg.org,
+                        project=cfg.project,
+                        domain=cfg.domain,
+                        name=name,
+                    )
+                )
+            )
+            found.add(name)
+        except Exception:
+            # Anything other than a hit is treated as "not there". A lookup failure
+            # only costs us a redundant create, which the server de-duplicates.
+            continue
+    return found