[Feat] add sleep plugins and examples (#3446)

Signed-off-by: machichima <nary12321@gmail.com>

Author: machichima

plugins/flytekit-sleep/examples/sleep_example.py

@@ -0,0 +1,49 @@
+"""
+Sleep Plugin Example
+====================
+
+The ``core-sleep`` plugin executes entirely in the backend — no task pod is created.
+The sleep duration is a normal task input, so it can be a dynamic workflow value.
+"""
+
+import os
+from datetime import timedelta
+
+from flytekitplugins.sleep import Sleep
+
+from flytekit import task, workflow
+
+
+@task(cache_version="2", task_config=Sleep())
+def sleep_for(duration: timedelta) -> None:
+    # This body only runs during local execution.
+    # On the cluster, the backend sleeps for `duration` without running this.
+    print(f"[local] sleeping for {duration}")
+
+
+@workflow
+def wf(duration: timedelta) -> None:
+    sleep_for(duration=duration)
+
+
+if __name__ == "__main__":
+    from click.testing import CliRunner
+
+    from flytekit.clis.sdk_in_container import pyflyte
+
+    runner = CliRunner()
+    path = os.path.realpath(__file__)
+    result = runner.invoke(
+        pyflyte.main,
+        [
+            "--config",
+            os.path.expanduser("~/.flyte/config-sandbox.yaml"),
+            "run",
+            "--remote",
+            path,
+            "wf",
+            "--duration",
+            "10s",
+        ],
+    )
+    print("Remote Execution: ", result.output)

plugins/flytekit-sleep/examples/sleep_fanout.py

@@ -0,0 +1,71 @@
+"""
+Sleep Fanout Example
+====================
+
+Fan out N core-sleep leaves in parallel using map_task.
+No task pods are created for the leaves — the backend handles sleep directly.
+
+Usage (remote):
+    python sleep_fanout.py
+"""
+
+import os
+from datetime import timedelta
+from typing import List
+
+from flytekitplugins.sleep import Sleep
+
+import flytekit as fl
+from flytekit import map_task, task, workflow
+
+sleep_image = fl.ImageSpec(
+    registry="ghcr.io/machichima",
+    name="sleep-fanout",
+    apt_packages=["git"],
+    packages=["git+https://github.com/machichima/flytekit.git@add-sleep-plugin#subdirectory=plugins/flytekit-sleep"],
+    env={"REBUILD": "1"},
+)
+
+
+@task(container_image=sleep_image)
+def make_durations(duration: timedelta, n: int) -> List[timedelta]:
+    return [duration] * n
+
+
+@task(task_config=Sleep(), container_image=sleep_image)
+def sleep_leaf(duration: timedelta) -> None:
+    pass
+
+
+@workflow
+def wf(sleep_duration: timedelta = timedelta(seconds=10), n_children: int = 400) -> None:
+    durations = make_durations(duration=sleep_duration, n=n_children)
+    map_task(sleep_leaf)(duration=durations)
+
+
+if __name__ == "__main__":
+    from click.testing import CliRunner
+
+    from flytekit.clis.sdk_in_container import pyflyte
+
+    runner = CliRunner()
+    path = os.path.realpath(__file__)
+
+    result = runner.invoke(
+        pyflyte.main,
+        [
+            "--config",
+            os.path.expanduser("~/.flyte/config-sandbox.yaml"),
+            "run",
+            "--remote",
+            path,
+            "wf",
+            "--sleep_duration",
+            "10s",
+            "--n_children",
+            "400",
+        ],
+    )
+    print(result.output)
+    if result.exception:
+        raise result.exception

plugins/flytekit-sleep/examples/sleep_fanout_dynamic.py

@@ -0,0 +1,76 @@
+"""
+Sleep Fanout (Dynamic) Example
+===============================
+
+Fan out N core-sleep leaves in parallel using @dynamic.
+n_children is a runtime input — the @dynamic task runs in a container to
+expand the subworkflow, then each leaf runs as core-sleep (no pod).
+
+Requires flytekitplugins-sleep in the container image. The ImageSpec below
+installs it from the local wheel via copy + pip install.
+
+Usage (remote):
+    python sleep_fanout_dynamic.py
+"""
+
+import os
+from datetime import timedelta
+
+from flytekitplugins.sleep import Sleep
+
+import flytekit as fl
+from flytekit import dynamic, task, workflow
+
+fanout_image = fl.ImageSpec(
+    registry="ghcr.io/machichima",
+    name="sleep-fanout",
+    apt_packages=["git"],
+    packages=["git+https://github.com/machichima/flytekit.git@add-sleep-plugin#subdirectory=plugins/flytekit-sleep"],
+)
+
+
+@task(task_config=Sleep())
+def sleep_leaf(duration: timedelta) -> None:
+    pass
+
+
+@dynamic(container_image=fanout_image)
+def sleep_fanout(n_children: int, sleep_duration: timedelta) -> None:
+    for _ in range(n_children):
+        sleep_leaf(duration=sleep_duration)
+
+
+@workflow
+def wf(
+    n_children: int = 10,
+    sleep_duration: timedelta = timedelta(seconds=10),
+) -> None:
+    sleep_fanout(n_children=n_children, sleep_duration=sleep_duration)
+
+
+if __name__ == "__main__":
+    from click.testing import CliRunner
+
+    from flytekit.clis.sdk_in_container import pyflyte
+
+    runner = CliRunner()
+    path = os.path.realpath(__file__)
+
+    result = runner.invoke(
+        pyflyte.main,
+        [
+            "--config",
+            os.path.expanduser("~/.flyte/config-sandbox.yaml"),
+            "run",
+            "--remote",
+            path,
+            "wf",
+            "--n_children",
+            "400",
+            "--sleep_duration",
+            "10s",
+        ],
+    )
+    print(result.output)
+    if result.exception:
+        raise result.exception

plugins/flytekit-sleep/flytekitplugins/__init__.py

@@ -0,0 +1,11 @@
+"""
+.. currentmodule:: flytekitplugins.sleep
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+   Sleep
+"""
+
+from .task import Sleep

plugins/flytekit-sleep/flytekitplugins/sleep/__init__.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Optional
+
+from flytekit.configuration import SerializationSettings
+from flytekit.core.python_function_task import PythonFunctionTask
+from flytekit.core.task import TaskPlugins
+
+
+@dataclass
+class Sleep:
+    """
+    Route a task to the backend ``core-sleep`` plugin.
+
+    The sleep duration is provided as a normal task input, not plugin config.
+    No container is launched; the backend handles the sleep directly.
+
+    Usage::
+
+        from flytekitplugins.sleep import Sleep
+        from flytekit import task
+        from datetime import timedelta
+
+        @task(task_config=Sleep())
+        def sleep_for(duration: timedelta) -> None:
+            pass  # only runs locally; backend executes the sleep
+    """
+
+
+class SleepFunctionTask(PythonFunctionTask[Sleep]):
+    _TASK_TYPE = "core-sleep"
+
+    def __init__(self, task_config: Optional[Sleep], task_function: Callable, **kwargs):
+        super().__init__(
+            task_config=task_config or Sleep(),
+            task_function=task_function,
+            task_type=self._TASK_TYPE,
+            **kwargs,
+        )
+
+    def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
+        return {}
+
+
+TaskPlugins.register_pythontask_plugin(Sleep, SleepFunctionTask)

plugins/flytekit-sleep/flytekitplugins/sleep/task.py

@@ -0,0 +1,36 @@
+from setuptools import setup
+
+PLUGIN_NAME = "sleep"
+
+microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
+
+__version__ = "0.0.0+develop"
+
+setup(
+    title="Sleep",
+    title_expanded="Flytekit Sleep Plugin",
+    name=microlib_name,
+    version=__version__,
+    author="flyteorg",
+    author_email="admin@flyte.org",
+    description="This package holds the core-sleep plugin for flytekit",
+    namespace_packages=["flytekitplugins"],
+    packages=[f"flytekitplugins.{PLUGIN_NAME}"],
+    install_requires=["flytekit>=1.15.1"],
+    license="apache2",
+    python_requires=">=3.9",
+    classifiers=[
+        "Intended Audience :: Science/Research",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: Apache Software License",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Software Development",
+        "Topic :: Software Development :: Libraries",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+    entry_points={"flytekit.plugins": [f"{PLUGIN_NAME}=flytekitplugins.{PLUGIN_NAME}"]},
+)

plugins/flytekit-sleep/setup.py

@@ -0,0 +1,69 @@
+from collections import OrderedDict
+from datetime import timedelta
+
+from flytekitplugins.sleep import Sleep
+from flytekitplugins.sleep.task import SleepFunctionTask
+
+from flytekit import task, workflow
+from flytekit.configuration import Image, ImageConfig, SerializationSettings
+from flytekit.extend import get_serializable
+
+default_img = Image(name="default", fqn="test", tag="tag")
+serialization_settings = SerializationSettings(
+    project="proj",
+    domain="dom",
+    version="123",
+    image_config=ImageConfig(default_image=default_img, images=[default_img]),
+    env={},
+)
+
+
+def test_task_type():
+    @task(task_config=Sleep())
+    def sleep_for(duration: timedelta) -> None:
+        pass
+
+    assert isinstance(sleep_for, SleepFunctionTask)
+    assert sleep_for.task_type == "core-sleep"
+
+
+def test_serialization():
+    @task(task_config=Sleep())
+    def sleep_for(duration: timedelta) -> None:
+        pass
+
+    task_spec = get_serializable(OrderedDict(), serialization_settings, sleep_for)
+
+    assert task_spec.template.type == "core-sleep"
+    assert task_spec.template.custom == {}
+
+    inputs = task_spec.template.interface.inputs
+    assert "duration" in inputs
+    assert len(inputs) == 1
+
+    outputs = task_spec.template.interface.outputs
+    assert len(outputs) == 0
+
+
+def test_local_execution_calls_function():
+    called = []
+
+    @task(task_config=Sleep())
+    def sleep_for(duration: timedelta) -> None:
+        called.append(duration)
+
+    sleep_for(duration=timedelta(seconds=1))
+    assert called == [timedelta(seconds=1)]
+
+
+def test_workflow_integration():
+    @task(task_config=Sleep())
+    def sleep_for(duration: timedelta) -> None:
+        pass
+
+    @workflow
+    def wf(duration: timedelta) -> None:
+        sleep_for(duration=duration)
+
+    spec = get_serializable(OrderedDict(), serialization_settings, wf)
+    assert len(spec.template.nodes) == 1