kfpint.py

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

"""
This file runs the KFP integration tests on KFP cluster. There's a number of
environment variables that need to be setup as well as the cluster.

See examples/pipelines/kfp/ for more information on how the cluster is used.

Cluster setup:

You'll need a KubeFlow Pipelines cluster with a torchserve instance with svc
name torchserve on the default namespace.

* https://www.kubeflow.org/docs/started/installing-kubeflow/
* https://github.com/pytorch/serve/blob/master/kubernetes/README.md

Environment variables:

```
export KFP_NAMESPACE=<kfp namespace>
export INTEGRATION_TEST_STORAGE=<cloud storage path>
export TORCHX_CONTAINER_REPO=<docker repo>
```

Once you have everything setup you can just run:

scripts/kfpint.py


"""

import argparse
import asyncio
import dataclasses
import json
import os
import os.path
import shutil
import subprocess
import tempfile
import time
from contextlib import contextmanager
from typing import Any, Callable, Iterator, Optional, TypeVar

import kfp

from integ_test_utils import (
    build_images,
    BuildInfo,
    getenv_asserts,
    MissingEnvError,
    push_images,
    run,
    run_in_bg,
)
from torchx.util.types import none_throws
from urllib3.exceptions import MaxRetryError

T = TypeVar("T")


def get_fn_name(fn: Callable[..., T]) -> str:
    if hasattr(fn, "__qualname__"):
        return fn.__qualname__
    elif hasattr(fn, "__name__"):
        return fn.__name__
    else:
        return "undefined"


def retry(f: Callable[..., T]) -> Callable[..., T]:
    retries: int = 5
    backoff: int = 3

    def wrapper(*args: Any, **kwargs: Any) -> T:
        curr_retries = 0
        while True:
            try:
                return f(*args, **kwargs)
            except:  # noqa: B001 E722
                if curr_retries == retries:
                    raise
                else:
                    sleep = backoff * 2**curr_retries
                    fn_name = get_fn_name(f)
                    print(f"retrying `{fn_name}` request after {sleep} seconds")
                    time.sleep(sleep)
                    curr_retries += 1
                    continue

    return wrapper


@retry
def get_client(host: str) -> kfp.Client:
    return kfp.Client(host=f"{host}/pipeline")


def get_free_port() -> int:
    return 32001


def enable_port_forward(local_port: int) -> "Optional[subprocess.Popen[str]]":
    # Enable port forward via running background process.
    # Kubernetes python does not support a clean way of
    # Kubernetes python cli provides a socket, more info:
    # https://github.com/kubernetes-client/python/blob/master/examples/pod_portforward.py
    # The drawback of this method is that we have to monkey patch
    # the urllib, which is used by the kfp client.
    # This approach is more cleaner than to use the python cli directly.
    try:
        namespace = getenv_asserts("KFP_NAMESPACE")
    except MissingEnvError:
        print("Skipping port forward due to workflow executed without env variable")
        return None
    return run_in_bg(
        "kubectl",
        "port-forward",
        "-n",
        namespace,
        "svc/ml-pipeline-ui",
        f"{local_port}:80",
    )


META_FILE = "meta"
IMAGES_FILE = "images.tar.zst"


def save_advanced_pipeline_spec(path: str, build: BuildInfo) -> None:
    print("generating advanced_pipeline spec")

    id = build.id
    torchx_image = build.torchx_image

    STORAGE_PATH = os.getenv("INTEGRATION_TEST_STORAGE", "/tmp/storage")
    root = os.path.join(STORAGE_PATH, id)
    output = os.path.join(root, "output")

    save_pipeline_spec(
        path,
        "advanced_pipeline.py",
        "--output_path",
        output,
        "--image",
        torchx_image,
        "--model_name",
        f"tiny_image_net_{id}",
    )


def save_pipeline_spec(path: str, pipeline_file: str, *args: str) -> None:
    print(f"generating pipeline spec for {pipeline_file}")

    with tempfile.TemporaryDirectory() as tmpdir:
        run(os.path.join("torchx/examples/pipelines/kfp", pipeline_file), *args)
        shutil.copy("pipeline.yaml", path)


@contextmanager
def path_or_tmp(path: Optional[str]) -> Iterator[str]:
    if path:
        os.makedirs(path, exist_ok=True)
        yield path
    else:
        with tempfile.TemporaryDirectory() as tmpdir:
            yield tmpdir


def _connection_error_message() -> str:
    kfp_host = getenv_asserts("KFP_HOST")
    return f"""
    Unable to connect to kfp cluster using {kfp_host}.
    Check that `kubectl` has proper credentials to execute port forward
    """


def save_build(path: str, build: BuildInfo) -> None:
    meta_path = os.path.join(path, META_FILE)
    with open(meta_path, "wt") as f:
        json.dump(dataclasses.asdict(build), f)


def run_pipeline(build: BuildInfo, pipeline_file: str) -> object:
    print(f"launching pipeline {pipeline_file}")
    HOST: str = getenv_asserts("KFP_HOST")

    try:
        client = get_client(HOST)
    except MaxRetryError:
        print(_connection_error_message())
        raise
    resp = client.create_run_from_pipeline_package(
        pipeline_file,
        arguments={},
        experiment_name="integration-tests",
        run_name=f"integration test {build.id} - {os.path.basename(pipeline_file)}",
    )
    ui_url = f"{HOST}/#/runs/details/{resp.run_id}"
    print(f"{resp.run_id} - launched! view run at {ui_url}")
    return resp


def wait_for_pipeline(
    resp: Any,  # pyre-fixme: KFP doesn't have a response type
) -> None:
    print(f"{resp.run_id} - waiting for completion")
    result = resp.wait_for_run_completion(
        timeout=1 * 60 * 60,
    )  # 1 hour
    print(f"{resp.run_id} - finished: {result}")
    assert result.run.status == "Succeeded", "run didn't succeed"


async def exec_job() -> None:
    parser = argparse.ArgumentParser(description="kfp integration test runner")
    parser.add_argument(
        "--path",
        type=str,
        help="path to place the files",
    )
    parser.add_argument(
        "--load",
        help="if specified load the build from path instead of building",
        action="store_true",
    )
    parser.add_argument(
        "--save",
        help="if specified save the build to path and exit",
        action="store_true",
    )
    parser.add_argument("--container_repo", type=str)
    args = parser.parse_args()

    with path_or_tmp(args.path) as path:
        advanced_pipeline_file = os.path.join(path, "advanced_pipeline.yaml")
        intro_pipeline_file = os.path.join(path, "intro_pipeline.yaml")
        dist_pipeline_file = os.path.join(path, "dist_pipeline.yaml")
        build = build_images()
        try:
            push_images(build, container_repo=args.container_repo)
        except MissingEnvError as e:
            print(f"Missing environments, only building: {e}")
            return
        finally:
            save_advanced_pipeline_spec(advanced_pipeline_file, build)
            save_pipeline_spec(intro_pipeline_file, "intro_pipeline.py")
            save_pipeline_spec(dist_pipeline_file, "dist_pipeline.py")

        pipeline_files = [
            advanced_pipeline_file,
            intro_pipeline_file,
            dist_pipeline_file,
        ]
        responses = [
            run_pipeline(build, pipeline_file) for pipeline_file in pipeline_files
        ]
        for response in responses:
            wait_for_pipeline(response)


def main() -> None:
    port = get_free_port()
    kfp_host = f"http://localhost:{port}"
    os.environ["KFP_HOST"] = kfp_host
    port_forward_proc = enable_port_forward(port)
    try:
        asyncio.run(exec_job())
    finally:
        if port_forward_proc:
            none_throws(port_forward_proc).kill()


if __name__ == "__main__":
    main()