import prefect
from prefect import task, Flow
from prefect.run_configs import KubernetesRun
from prefect.storage.gcs import GCS

from prefect_flows.utils.run_config import make_kubernetes_run_config

@task
def hello_task():
    logger = prefect.context.get("logger")
    <http://logger.info|logger.info>("Hello world!")

with Flow(FLOW_NAME) as flow:
    hello_task()

flow.run_config = KubernetesRun(
    job_template=make_kubernetes_run_config(
        run_name="example02",
        prefect_cloud_auth_token=<token>,
    ),
)

flow.storage = GCS(
    project=<project>,
    bucket=f"well-{<env>}-prefect-flow-mlops",
    key="flows/example/example02.py",
)

import uuid


def make_kubernetes_run_config(run_name: str, prefect_cloud_auth_token: str) -> str:
  run_id = uuid.uuid4()
  return f"""
apiVersion: batch/v1
kind: Job
metadata:
  name: prefect-job-{run_name}-{run_id}
  namespace: mlops-prefect-agent
  labels:
    app: prefect-job-{run_name}-{run_id}
    identifier: {run_name}-{run_id}
spec:
  template:
    metadata:
      labels:
        app: prefect-job-{run_name}-{run_id}
        identifier: {run_name}-{run_id}
    spec:
      containers:
        - name: flow
          serviceAccountName: prefect-agent
          image: prefecthq/prefect:latest-python3.8
          imagePullPolicy: Always
          command: ["/bin/sh", "-c"]
          args: ["prefect execute flow-run"]
          env:
            - name: PREFECT__CLOUD__API
              value: <https://api.prefect.io>
            - name: PREFECT__CLOUD__AUTH_TOKEN
              value: {prefect_cloud_auth_token}
            - name: PREFECT__CONTEXT__FLOW_RUN_ID
              value: ""
            - name: PREFECT__CONTEXT__NAMESPACE
              value: mlops-prefect-agent
            - name: PREFECT__CLOUD__USE_LOCAL_SECRETS
              value: "false"
            - name: PREFECT__LOGGING__LOG_TO_CLOUD
              value: "true"
            - name: PREFECT__LOGGING__LEVEL
              value: "DEBUG"
            - name: PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS
              value: "prefect.engine.cloud.CloudFlowRunner"
            - name: PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS
              value: "prefect.engine.cloud.CloudTaskRunner"
          resources:
            requests:
              cpu: "200m"
            limits:
              cpu: "200m"
      restartPolicy: Never
      imagePullSecrets:
      - name: ""
"""

[2021-09-26 18:36:49,279] INFO - agent | Deploying flow run 1cb8289f-932d-4ae6-823d-a4fdfecacbde to execution environment...
[2021-09-26 18:36:49,606] INFO - agent | Completed deployment of flow run 1cb8289f-932d-4ae6-823d-a4fdfecacbde
INFO:agent:Completed deployment of flow run 1cb8289f-932d-4ae6-823d-a4fdfecacbde