from dask_kubernetes.operator import KubeCluster
from prefect import task, flow
from prefect_dask import DaskTaskRunner
from prefect_aws.s3 import S3Bucket
from prefect.infrastructure.kubernetes import KubernetesJob

import logging
logging.basicConfig(encoding='utf-8', level=logging.DEBUG)

logger = logging.getLogger(__name__)

task_runner = DaskTaskRunner(
    cluster_class=KubeCluster,
    cluster_kwargs=dict(
        name='prefect_dask',
        image='<http://ghcr.io/dask/dask:latest|ghcr.io/dask/dask:latest>',
        n_workers=1,
    ),
)

# Define some tasks for us to run in our flow

@task
def extract() -> list:
    return [1, 2, 3, 4, 5, 6]

@task(retries=3, retry_delay_seconds=15)
def transform(number: int) -> int:
    return number * 2

@task()
def load(numbers: list) -> list:
    return [i for i in numbers if i]

@flow(
    name="Dask Kubernetes Flow",
    task_runner=task_runner,
)
def main():
    numbers = extract()
    tranformed_numbers = transform.map(numbers)
    numbers_twice = transform.map(tranformed_numbers)
    result = load(numbers=numbers_twice)

from prefect.filesystems import S3
from prefect.deployments import Deployment
from dask_kubernetes.operator import KubeCluster
from prefect import task, flow
from prefect_dask import DaskTaskRunner
from prefect_aws.s3 import S3Bucket
from prefect.infrastructure.kubernetes import KubernetesJob
import ephemeral_dask

import logging
logging.basicConfig(encoding='utf-8', level=logging.DEBUG)

logger = logging.getLogger(__name__)

storage = S3Bucket.load("k8s-minio")  # load a pre-defined block
infra = KubernetesJob.load('k8s')

deployment = Deployment.build_from_flow(
    flow=ephemeral_dask.main,
    name="k8s-test",
    work_queue_name="kubernetes",
    storage=storage,
    infrastructure=infra,
)

deployment.apply()

from prefect.infrastructure.kubernetes import KubernetesJob, KubernetesImagePullPolicy
from prefect_aws.s3 import S3Bucket, MinIOCredentials

minio_cred = MinIOCredentials(
    minio_root_user='minioadmin',
    minio_root_password='minioadmin'
)
minio_cred.save('minio-cred', overwrite=True)

bucket = S3Bucket(
    bucket_name='k8s',
    minio_credentials=minio_cred,
    endpoint_url="<https://minio:9000>",
)
bucket.save('k8s-minio', overwrite=True)

k8_job = KubernetesJob(
    image_pull_policy=KubernetesImagePullPolicy.ALWAYS,
    env={
        "EXTRA_PIP_PACKAGES": "prefect_aws",
    },
)
k8_job.save('k8s', overwrite=True)

apiVersion: <http://rbac.authorization.k8s.io/v1|rbac.authorization.k8s.io/v1>
kind: RoleBinding
metadata:
  name: default
  namespace: default
subjects:

# Reference to upper's `metadata.name`
    name: default
    # Reference to upper's `metadata.namespace`
    namespace: default
roleRef:
  kind: ClusterRole
  name: cluster-admin
  apiGroup: <http://rbac.authorization.k8s.io|rbac.authorization.k8s.io>