Add the Flyte agent to provision and manage K8s (data) service for deep learning (GNN) use cases (#3004)

Signed-off-by: Shuying Liang <shuying.liang@gmail.com>
Signed-off-by: Future-Outlier <eric901201@gmail.com>
Signed-off-by: JiaWei Jiang <waynechuang97@gmail.com>
Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>
Co-authored-by: Han-Ru Chen (Future-Outlier) <eric901201@gmail.com>
Co-authored-by: Thomas J. Fan <thomasjpfan@gmail.com>
Co-authored-by: 江家瑋 <36886416+JiangJiaWei1103@users.noreply.github.com>
Co-authored-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Co-authored-by: Niels Bantilan <niels.bantilan@gmail.com>

Author: 6 people

.github/workflows/pythonbuild.yml

@@ -303,6 +303,7 @@ jobs:
           - flytekit-huggingface
           - flytekit-identity-aware-proxy
           - flytekit-inference
+          - flytekit-k8sdataservice
           - flytekit-k8s-pod
           - flytekit-kf-mpi
           - flytekit-kf-pytorch

.pre-commit-config.yaml

@@ -1,7 +1,7 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.6.9
+    rev: v0.8.3
     hooks:
       # Run the linter.
       - id: ruff
@@ -26,5 +26,6 @@ repos:
     rev: v2.3.0
     hooks:
       - id: codespell
-        additional_dependencies:
-          - tomli
+        args:
+          - --ignore-words-list=assertIn  # Ignore 'assertIn'
+        additional_dependencies: [tomli]

Dockerfile.agent

@@ -11,6 +11,7 @@ RUN apt-get update && apt-get install build-essential -y \
 RUN uv pip install --system --no-cache-dir -U flytekit==$VERSION \
   flytekitplugins-airflow==$VERSION \
   flytekitplugins-bigquery==$VERSION \
+  flytekitplugins-k8sdataservice==$VERSION \
   flytekitplugins-openai==$VERSION \
   flytekitplugins-snowflake==$VERSION \
   flytekitplugins-awssagemaker==$VERSION \

docs/source/plugins/k8sstatefuldataservice.rst

@@ -0,0 +1,12 @@
+.. k8sstatefuldataservice:
+
+###################################################
+Kubernetes StatefulSet Data Service API reference
+###################################################
+
+.. tags:: Integration, DeepLearning, MachineLearning, Kubernetes, GNN
+
+.. automodule:: flytekitplugins.k8sdataservice
+   :no-members:
+   :no-inherited-members:
+   :no-special-members:

plugins/flytekit-k8sdataservice/README.md

@@ -0,0 +1,104 @@
+# K8s Stateful Service Plugin
+
+This plugin provides support for Kubernetes StatefulSet and Service integration, enabling seamless provisioning and coordination with any Kubernetes services or Flyte tasks. It is especially suited for deep learning use cases at scale, where distributed and parallelized data loading and caching across nodes are required.
+
+## Features
+- **Predictable and Reliable Endpoints**: The service creates consistent endpoints, facilitating communication between services or tasks within the same Kubernetes cluster.
+- **Reusable Across Runs**: Service tasks can persist across task runs, ensuring consistency. Alternatively, a cleanup sensor can release cluster resources when they are no longer needed.
+- **Conventional Pod Naming**: Pods in the StatefulSet follow a conventional naming pattern. For instance, if the StatefulSet name is `foo` and replicas are set to 2, the pod endpoints will be `foo-0.foo:1234` and `foo-1.foo:1234`. This simplifies endpoint construction for training or inference scripts. For example, gRPC endpoints can directly use `foo-0.foo:1234` and `foo-1.foo:1234`.
+
+## Installation
+
+Install the plugin via pip:
+
+```bash
+pip install flytekitplugins-k8sdataservice
+```
+
+## Usage
+
+Below is an example demonstrating how to provision and run a service in Kubernetes, making it reachable within the cluster.
+
+**Note**: Utility functions are available to generate unique service names that can be reused across training or inference scripts.
+
+### Example Usage
+
+#### Provisioning a Data Service
+```python
+from flytekitplugins.k8sdataservice import DataServiceConfig, DataServiceTask, CleanupSensor
+from utils.infra import gen_infra_name
+from flytekit import kwtypes, Resources, task, workflow
+
+# Generate a unique infrastructure name
+name = gen_infra_name()
+
+def k8s_data_service():
+    gnn_config = DataServiceConfig(
+        Name=name,
+        Requests=Resources(cpu='1', mem='1Gi'),
+        Limits=Resources(cpu='2', mem='2Gi'),
+        Replicas=1,
+        Image="busybox:latest",
+        Command=[
+            "bash",
+            "-c",
+            "echo Hello Flyte K8s Stateful Service! && sleep 3600"
+        ],
+    )
+
+    gnn_task = DataServiceTask(
+        name="K8s Stateful Data Service",
+        inputs=kwtypes(ds=str),
+        task_config=gnn_config,
+    )
+    return gnn_task
+
+# Define a cleanup sensor
+gnn_sensor = CleanupSensor(name="Cleanup")
+
+# Define a workflow to test the data service
+@workflow
+def test_dataservice_wf(name: str):
+    k8s_data_service()(ds="OSS Flyte K8s Data Service Demo") \
+          >> gnn_sensor(
+              release_name=name,
+              cleanup_data_service=True,
+          )
+
+if __name__ == "__main__":
+    out = test_dataservice_wf(name="example")
+    print(f"Running test_dataservice_wf() {out}")
+```
+
+#### Accessing the Data Service
+Other tasks or services that need to access the data service can do so in multiple ways. For example, using environment variables:
+
+```python
+from kubernetes.client import V1PodSpec, V1Container, V1EnvVar
+
+PRIMARY_CONTAINER_NAME = "primary"
+FLYTE_POD_SPEC = V1PodSpec(
+    containers=[
+        V1Container(
+            name=PRIMARY_CONTAINER_NAME,
+            env=[
+                V1EnvVar(name="MY_DATASERVICES", value=f"{name}-0.{name}:40000 {name}-1.{name}:40000"),
+            ],
+        )
+    ],
+)
+
+task_config = MPIJob(
+    launcher=Launcher(replicas=1, pod_template=FLYTE_POD_SPEC),
+    worker=Worker(replicas=1, pod_template=FLYTE_POD_SPEC),
+)
+
+@task(task_config=task_config)
+def mpi_task() -> str:
+    return "your script uses the envs to communicate with the data service "
+```
+
+### Key Points
+- The `DataServiceConfig` defines resource requests, limits, replicas, and the container image/command.
+- The `CleanupSensor` ensures resources are cleaned up when required.
+- The workflow connects the service provisioning and cleanup process for streamlined operations.

plugins/flytekit-k8sdataservice/dev-requirements.txt

@@ -0,0 +1 @@
+kubernetes~=23.6.0

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

@@ -0,0 +1,15 @@
+"""
+.. currentmodule:: flytekitplugins.k8sdataservice
+
+This package contains things that are useful when extending Flytekit.
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+   DataServiceTask
+"""
+
+from .agent import DataServiceAgent  # noqa: F401
+from .sensor import CleanupSensor  # noqa: F401
+from .task import DataServiceConfig, DataServiceTask  # noqa: F401

plugins/flytekit-k8sdataservice/flytekitplugins/k8sdataservice/agent.py

@@ -0,0 +1,98 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from flyteidl.core.execution_pb2 import TaskExecution
+from flytekitplugins.k8sdataservice.k8s.manager import K8sManager
+from flytekitplugins.k8sdataservice.task import DataServiceConfig
+
+from flytekit import logger
+from flytekit.extend.backend.base_agent import AgentRegistry, AsyncAgentBase, Resource, ResourceMeta
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import TaskTemplate
+
+
+@dataclass
+class DataServiceMetadata(ResourceMeta):
+    dataservice_config: DataServiceConfig
+    name: str
+
+
+class DataServiceAgent(AsyncAgentBase):
+    name = "K8s DataService Async Agent"
+
+    def __init__(self):
+        self.k8s_manager = K8sManager()
+        super().__init__(task_type_name="dataservicetask", metadata_type=DataServiceMetadata)
+        self.config = None
+
+    def create(
+        self, task_template: TaskTemplate, output_prefix: str, inputs: Optional[LiteralMap] = None, **kwargs
+    ) -> DataServiceMetadata:
+        graph_engine_config = task_template.custom
+        self.k8s_manager.set_configs(graph_engine_config)
+        logger.info(f"Loaded agent config file {self.config}")
+        existing_release_name = graph_engine_config.get("ExistingReleaseName", None)
+        logger.info(f"The existing data service release name is {existing_release_name}")
+
+        name = ""
+        if existing_release_name is None or existing_release_name == "":
+            logger.info("Creating K8s data service resources...")
+            name = self.k8s_manager.create_data_service()
+            logger.info(f'Data service {name} with image {graph_engine_config["Image"]} completed')
+        else:
+            name = existing_release_name
+            logger.info(f"User configs to use the existing data service release name: {name}.")
+
+        dataservice_config = DataServiceConfig(
+            Name=graph_engine_config.get("Name", None),
+            Image=graph_engine_config["Image"],
+            Command=graph_engine_config["Command"],
+            Cluster=graph_engine_config["Cluster"],
+            ExistingReleaseName=graph_engine_config.get("ExistingReleaseName", None),
+        )
+        metadata = DataServiceMetadata(
+            dataservice_config=dataservice_config,
+            name=name,
+        )
+        logger.info(f"Created DataService metadata {metadata}")
+        return metadata
+
+    def get(self, resource_meta: DataServiceMetadata) -> Resource:
+        logger.info("K8s Data Service get is called")
+        data = resource_meta.dataservice_config
+        data_dict = data.__dict__ if isinstance(data, DataServiceConfig) else data
+        logger.info(f"The data_dict is {data_dict}")
+        self.k8s_manager.set_configs(data_dict)
+        name = data.Name
+        logger.info(f"Get the stateful set name {name}")
+
+        k8s_status = self.k8s_manager.check_stateful_set_status(name)
+        flyte_state = None
+        if k8s_status in ["failed", "timeout", "timedout", "canceled", "skipped", "internal_error"]:
+            flyte_state = TaskExecution.FAILED
+        elif k8s_status in ["done", "succeeded", "success"]:
+            flyte_state = TaskExecution.SUCCEEDED
+        elif k8s_status in ["running", "terminating", "pending"]:
+            flyte_state = TaskExecution.RUNNING
+        else:
+            logger.error(f"Unrecognized state: {k8s_status}")
+        outputs = {
+            "data_service_name": name,
+        }
+        # TODO: Add logs for StatefulSet.
+        return Resource(phase=flyte_state, outputs=outputs)
+
+    def delete(self, resource_meta: DataServiceMetadata):
+        logger.info("DataService delete is called")
+        data = resource_meta.dataservice_config
+
+        data_dict = data.__dict__ if isinstance(data, DataServiceConfig) else data
+        self.k8s_manager.set_configs(data_dict)
+
+        name = resource_meta.name
+        logger.info(f"To delete the DataService (e.g., StatefulSet and Service) with name {name}")
+        self.k8s_manager.delete_stateful_set(name)
+        self.k8s_manager.delete_service(name)
+
+
+AgentRegistry.register(DataServiceAgent())

plugins/flytekit-k8sdataservice/flytekitplugins/k8sdataservice/k8s/__init__.py

@@ -0,0 +1,20 @@
+from kubernetes import config
+
+from flytekit import logger
+
+
+class KubeConfig:
+    def __init__(self):
+        pass
+
+    def load_kube_config(self) -> None:
+        """Load the kubernetes config based on fabric details prior to K8s client usage
+
+        :params target_fabric: fabric on which we are loading configs
+        """
+        try:
+            logger.info("Attempting to load in-cluster configuration.")
+            config.load_incluster_config()  # This will use the service account credentials
+            logger.info("Successfully loaded in-cluster configuration using the agent service account.")
+        except config.ConfigException as e:
+            logger.warning(f"Failed to load in-cluster configuration. {e}")