Prefect flow with KubeCluster, GCS storage, and KubernetesRun

test_flow_kubecluster_gcs_storage.py

import urllib.request

import prefect
from dask_kubernetes import KubeCluster, make_pod_spec
from prefect import task, unmapped, Flow
from prefect.executors import DaskExecutor
from prefect.run_configs import KubernetesRun
from prefect.storage import GCS


@task
def get_html(url):
    logger = prefect.context.get("logger")
    logger.info(f"Downloading {url}")
    with urllib.request.urlopen(url) as response:
        return (url, response.read())


@task
def write_html(url_html):
    logger = prefect.context.get("logger")
    url, html = url_html
    basename = f"{url.split('://')[-1].strip('/').replace('/', '_')}.html"
    outname = f"/tmp/{basename}"
    logger.info(f"Writing {outname}")
    with open(outname, "wb") as fh:
        fh.write(html)


with Flow("Basic web scrape with KubeCluster") as flow:
    urls = ["https://example.com", "https://docs.python.org/3/howto/urllib2.html"]
    payload = get_html.map(urls)
    write_html.map(payload)

pod_spec = make_pod_spec(
    image="us-west2-docker.pkg.dev/hai-gcp-augmenting/agena-watch-docker-images/prefect-agenda-watch:latest",
    env={
        "GOOGLE_APPLICATION_CREDENTIALS": "/tmp/gcp_secrets.json",
    },
)
executor = DaskExecutor(
    cluster_class=KubeCluster, cluster_kwargs={"pod_template": pod_spec}
)
flow.storage = GCS(
    bucket="aw-prefect-flows",
    project="hai-gcp-augmenting",
)
flow.run_config = KubernetesRun(
    env={"GOOGLE_APPLICATION_CREDENTIALS": "/tmp/gcp_secrets.json"},
    image="us-west2-docker.pkg.dev/hai-gcp-augmenting/agena-watch-docker-images/prefect-agenda-watch:latest",
)
flow.run(executor=executor)