text: <<<<<<< HEAD
  from collections import defaultdict
  from datetime import datetime
  import os
  import pytz
  import time
  import uuid
  from typing import Optional, Iterable, List, Any
  =======
  import os
  import time
  import uuid
  from typing import Iterable, List, Any
  >>>>>>> prefect clone
  
  import json
  import yaml
  
  import prefect
  from prefect import config
  from prefect.agent import Agent
  from prefect.engine.state import Failed
  <<<<<<< HEAD
  from prefect.run_configs import KubernetesRun
  from prefect.utilities.agent import get_flow_image, get_flow_run_command
  from prefect.utilities.exceptions import ClientError
  =======
  from prefect.serialization.run_config import RunConfigSchema
  from prefect.utilities.agent import get_flow_image, get_flow_run_command
  >>>>>>> prefect clone
  from prefect.utilities.filesystems import read_bytes_from_path
  from prefect.utilities.graphql import GraphQLResult
  
  DEFAULT_JOB_TEMPLATE_PATH = os.path.join(os.path.dirname(__file__), "job_template.yaml")
  
  
  def _get_or_create(d: dict, key: str, val: Any = None) -> Any:
      """Get a (possibly nested) field from a dict, creating intermediate values
      if needed."""
      if val is None:
          val = {}
      path = key.split(".")
      for k in path[:-1]:
          d = d.setdefault(k, {})
      return d.setdefault(path[-1], val)
  
  
  class KubernetesAgent(Agent):
      """
      Agent which deploys flow runs as Kubernetes jobs. Currently this is required to either
      run on a k8s cluster or on a local machine where the kube_config is pointing at the
      desired cluster. Information on using the Kubernetes Agent can be found at
      https://docs.prefect.io/orchestration/agents/kubernetes.html
  
  <<<<<<< HEAD
  =======
      Environment variables may be set on the agent to be provided to each flow run's job:
      ```
      prefect agent kubernetes start --env MY_SECRET_KEY=secret --env OTHER_VAR=$OTHER_VAR
      ```
  
      These can also be used to control the k8s job spec that describes the flow run jobs.
      For example, to set the k8s secret used to pull images from a non-public registry:
      ```
      prefect agent kubernetes start --env IMAGE_PULL_SECRETS=my-img-pull-secret
      ```
  
      For details on the available environment variables for customizing the job spec,
      see `help(KubernetesAgent.generate_job_spec_from_environment)`.
  
      Specifying a namespace for the agent will create flow run jobs in that namespace:
      ```
      prefect agent kubernetes start --namespace dev
      ```
  
  >>>>>>> prefect clone
      Args:
          - agent_config_id (str, optional): An optional agent configuration ID that can be used to set
              configuration based on an agent from a backend API. If set all configuration values will be
              pulled from backend agent configuration.
          - namespace (str, optional): A Kubernetes namespace to create jobs in. Defaults
              to the environment variable `NAMESPACE` or `default`.
  <<<<<<< HEAD
          - service_account_name (str, optional): A kubernetes service account name to use by
              default for created jobs. May be overridden per-flow by specifying
              on a flow's `KubernetesRun` run config.
          - image_pull_secrets (list, optional): A list of image pull secrets to use by default
              for created jobs. May be overridden per-flow by specifying on a flow's
              `KubernetesRun` run config.
  =======
  >>>>>>> prefect clone
          - job_template_path (str, optional): A path to a job template file to use instead
              of the default.
          - name (str, optional): An optional name to give this agent. Can also be set through
              the environment variable `PREFECT__CLOUD__AGENT__NAME`. Defaults to "agent"
          - labels (List[str], optional): a list of labels, which are arbitrary string
              identifiers used by Prefect Agents when polling for work
          - env_vars (dict, optional): a dictionary of environment variables and values that will
              be set on each flow run that this agent submits for execution and in the agent's
              own environment
          - max_polls (int, optional): maximum number of times the agent will poll Prefect Cloud
              for flow runs; defaults to infinite
          - agent_address (str, optional):  Address to serve internal API at. Currently this is
              just health checks for use by an orchestration layer. Leave blank for no API server
              (default).
          - no_cloud_logs (bool, optional): Disable logging to a Prefect backend for this agent
              and all deployed flow runs
          - volume_mounts (list, optional): A list of volumeMounts to mount when a job is
              run. The volumeMounts in the list should be specified as dicts
              i.e `[{"name": "my-vol", "mountPath": "/mnt/my-mount"}]`
          - volumes (list, optional): A list of volumes to make available to be mounted when a
              job is run. The volumes in the list should be specified as nested dicts.
              i.e `[{"name": "my-vol", "csi": {"driver": "secrets-store.csi.k8s.io"}}]`
      """
  
      def __init__(
          self,
          agent_config_id: str = None,
          namespace: str = None,
  <<<<<<< HEAD
          service_account_name: str = None,
          image_pull_secrets: Iterable[str] = None,
  =======
  >>>>>>> prefect clone
          job_template_path: str = None,
          name: str = None,
          labels: Iterable[str] = None,
          env_vars: dict = None,
          max_polls: int = None,
          agent_address: str = None,
          no_cloud_logs: bool = False,
          volume_mounts: List[dict] = None,
          volumes: List[dict] = None,
      ) -> None:
          super().__init__(
              agent_config_id=agent_config_id,
              name=name,
              labels=labels,
              env_vars=env_vars,
              max_polls=max_polls,
              agent_address=agent_address,
              no_cloud_logs=no_cloud_logs,
          )
  
          self.namespace = namespace or os.getenv("NAMESPACE", "default")
  <<<<<<< HEAD
          self.service_account_name = service_account_name or os.getenv(
              "SERVICE_ACCOUNT_NAME"
          )
          if image_pull_secrets is None:
              image_pull_secrets_env = os.getenv("IMAGE_PULL_SECRETS")
              image_pull_secrets = (
                  [s.strip() for s in image_pull_secrets_env.split(",")]
                  if image_pull_secrets_env is not None
                  else None
              )
          self.image_pull_secrets = image_pull_secrets
  =======
  >>>>>>> prefect clone
          self.job_template_path = job_template_path or DEFAULT_JOB_TEMPLATE_PATH
          self.volume_mounts = volume_mounts
          self.volumes = volumes
  
          from kubernetes import client, config
  
          try:
              self.logger.debug("Loading incluster configuration")
              config.load_incluster_config()
          except config.config_exception.ConfigException as exc:
              self.logger.warning(
                  "{} Using out of cluster configuration option.".format(exc)
              )
              self.logger.debug("Loading out of cluster configuration")
              config.load_kube_config()
  
          self.batch_client = client.BatchV1Api()
          self.core_client = client.CoreV1Api()
          self.k8s_client = client
  
  <<<<<<< HEAD
          min_datetime = datetime.min.replace(tzinfo=pytz.UTC)
          self.job_pod_event_timestamps = defaultdict(  # type: ignore
              lambda: defaultdict(lambda: min_datetime)
          )
  
  =======
  >>>>>>> prefect clone
          self.logger.debug(f"Namespace: {self.namespace}")
  
      def manage_jobs(self) -> None:
          """
          This function checks if jobs are `Failed` or `Succeeded` and if they are then the jobs are
          deleted from the namespace. If one of the job's pods happen to run into image pulling errors
          then the flow run is failed and the job is still deleted.
          """
          self.logger.debug(
              "Retrieving information of jobs that are currently in the cluster..."
          )
  
          more = True
          _continue = ""
          while more:
              try:
                  jobs = self.batch_client.list_namespaced_job(
                      namespace=self.namespace,
                      label_selector="prefect.io/identifier",
                      limit=20,
                      _continue=_continue,
                  )
                  _continue = jobs.metadata._continue
                  more = bool(_continue)
  
                  for job in jobs.items:
                      delete_job = job.status.failed or job.status.succeeded
                      job_name = job.metadata.name
                      flow_run_id = job.metadata.labels.get("prefect.io/flow_run_id")
  
  <<<<<<< HEAD
                      # Check for pods that are stuck with image pull errors
  =======
  >>>>>>> prefect clone
                      if not delete_job:
                          pods = self.core_client.list_namespaced_pod(
                              namespace=self.namespace,
                              label_selector="prefect.io/identifier={}".format(
                                  job.metadata.labels.get("prefect.io/identifier")
                              ),
                          )
  
                          for pod in pods.items:
  <<<<<<< HEAD
                              pod_name = pod.metadata.name
  =======
  >>>>>>> prefect clone
                              if pod.status.container_statuses:
                                  for container_status in pod.status.container_statuses:
                                      waiting = container_status.state.waiting
                                      if waiting and (
                                          waiting.reason == "ErrImagePull"
                                          or waiting.reason == "ImagePullBackOff"
                                      ):
                                          self.logger.debug(
                                              f"Failing flow run {flow_run_id} due to pod {waiting.reason}"
                                          )
  <<<<<<< HEAD
                                          try:
                                              self.client.set_flow_run_state(
                                                  flow_run_id=flow_run_id,
                                                  state=Failed(
                                                      message="Kubernetes Error: {}".format(
                                                          container_status.state.waiting.message
                                                      )
                                                  ),
                                              )
                                          except ClientError as exc:
                                              self.logger.error(
                                                  "Error attempting to set flow run state for "
                                                  f"{flow_run_id}: "
                                                  f"{exc}"
                                              )
  =======
                                          self.client.set_flow_run_state(
                                              flow_run_id=flow_run_id,
                                              state=Failed(
                                                  message="Kubernetes Error: {}".format(
                                                      container_status.state.waiting.message
                                                  )
                                              ),
                                          )
  >>>>>>> prefect clone
  
                                          delete_job = True
                                          break
  
  <<<<<<< HEAD
                              # Report recent events for pending pods to flow run logs
                              if pod.status.phase == "Pending":
                                  pod_events = self.core_client.list_namespaced_event(
                                      namespace=self.namespace,
                                      field_selector="involvedObject.name={}".format(
                                          pod_name
                                      ),
                                      timeout_seconds=30,
                                  )
  
                                  for event in sorted(
                                      pod_events.items, key=lambda x: x.last_timestamp
                                  ):
                                      # Skip old events
                                      if (
                                          event.last_timestamp
                                          < self.job_pod_event_timestamps[job_name][
                                              pod_name
                                          ]
                                      ):
                                          continue
  
                                      self.job_pod_event_timestamps[job_name][
                                          pod_name
                                      ] = event.last_timestamp
  
                                      log_msg = (
                                          f"Event: {event.reason!r} on pod {pod_name!r}\n"
                                          f"\tMessage: {event.message}"
                                      )
  
                                      # Send pod failure information to flow run logs
                                      self.client.write_run_logs(
                                          [
                                              dict(
                                                  flow_run_id=flow_run_id,
                                                  name="k8s-infra",
                                                  message=log_msg,
                                                  level="DEBUG",
                                                  timestamp=event.last_timestamp.isoformat(),
                                              )
                                          ]
                                      )
  
                      # Report failed pods
                      if job.status.failed:
                          pods = self.core_client.list_namespaced_pod(
                              namespace=self.namespace,
                              label_selector="prefect.io/identifier={}".format(
                                  job.metadata.labels.get("prefect.io/identifier")
                              ),
                          )
  
                          failed_pods = []
                          for pod in pods.items:
                              if pod.status.phase != "Failed":
                                  continue
  
                              # Format pod failure error message
                              failed_pods.append(pod.metadata.name)
                              pod_status_logs = [f"Pod {pod.metadata.name} failed."]
                              if not pod.status.container_statuses:
                                  pod_status_logs.append(
                                      "\tNo container statuses found for pod"
                                  )
                              else:
                                  for status in pod.status.container_statuses:
                                      state = (
                                          "running"
                                          if status.state.running
                                          else "waiting"
                                          if status.state.waiting
                                          else "terminated"
                                          if status.state.terminated
                                          else "Not Found"
                                      )
                                      pod_status_logs.append(
                                          f"\tContainer '{status.name}' state: {state}"
                                      )
  
                                      if status.state.terminated:
                                          pod_status_logs.append(
                                              f"\t\tExit Code:: {status.state.terminated.exit_code}"
                                          )
                                          if status.state.terminated.message:
                                              pod_status_logs.append(
                                                  f"\t\tMessage: {status.state.terminated.message}"
                                              )
                                          if status.state.terminated.reason:
                                              pod_status_logs.append(
                                                  f"\t\tReason: {status.state.terminated.reason}"
                                              )
                                          if status.state.terminated.signal:
                                              pod_status_logs.append(
                                                  f"\t\tSignal: {status.state.terminated.signal}"
                                              )
  
                              # Send pod failure information to flow run logs
                              self.client.write_run_logs(
                                  [
                                      dict(
                                          flow_run_id=flow_run_id,
                                          name="k8s-infra",
                                          message="\n".join(pod_status_logs),
                                          level="ERROR",
                                      )
                                  ]
                              )
  
                          # If there are failed pods and the run is not finished, fail the run
                          if (
                              failed_pods
                              and not self.client.get_flow_run_state(
                                  flow_run_id
                              ).is_finished()
                          ):
                              self.logger.debug(
                                  f"Failing flow run {flow_run_id} due to the failed pods {failed_pods}"
                              )
                              try:
                                  self.client.set_flow_run_state(
                                      flow_run_id=flow_run_id,
                                      state=Failed(
                                          message="Kubernetes Error: pods {} failed for this job".format(
                                              failed_pods
                                          )
                                      ),
                                  )
                              except ClientError as exc:
                                  self.logger.error(
                                      f"Error attempting to set flow run state for {flow_run_id}: "
                                      f"{exc}"
                                  )
  
                      # Delete job if it is successful or failed
                      if delete_job:
                          self.logger.debug(f"Deleting job {job_name}")
                          try:
                              self.job_pod_event_timestamps.pop(job_name, None)
  =======
                      if delete_job:
                          self.logger.debug(f"Deleting job {job_name}")
                          try:
  >>>>>>> prefect clone
                              self.batch_client.delete_namespaced_job(
                                  name=job_name,
                                  namespace=self.namespace,
                                  body=self.k8s_client.V1DeleteOptions(
                                      propagation_policy="Foreground"
                                  ),
                              )
                          except self.k8s_client.rest.ApiException as exc:
                              if exc.status != 404:
                                  self.logger.error(
                                      f"{exc.status} error attempting to delete job {job_name}"
                                  )
              except self.k8s_client.rest.ApiException as exc:
                  if exc.status == 410:
                      self.logger.debug("Refreshing job listing token...")
                      _continue = ""
                      continue
                  else:
                      self.logger.debug(exc)
  
      def heartbeat(self) -> None:
          """
          Check status of jobs created by this agent, delete completed jobs and failed containers.
          """
          try:
              self.manage_jobs()
          except Exception:
              self.logger.error("Error while managing existing k8s jobs", exc_info=True)
          super().heartbeat()
  
      def deploy_flow(self, flow_run: GraphQLResult) -> str:
          """
          Deploy flow runs on to a k8s cluster as jobs
  
          Args:
              - flow_run (GraphQLResult): A GraphQLResult flow run object
  
          Returns:
              - str: Information about the deployment
          """
          import urllib3.exceptions
  
          self.logger.info("Deploying flow run {}".format(flow_run.id))  # type: ignore
  
          job_spec = self.generate_job_spec(flow_run=flow_run)
          job_name = job_spec["metadata"]["name"]
  
          self.logger.debug("Creating namespaced job {}".format(job_name))
          attempts = 3
          while True:
              try:
                  self.batch_client.create_namespaced_job(
                      namespace=self.namespace, body=job_spec
                  )
                  break
              except self.k8s_client.rest.ApiException as exc:
                  if exc.status == 409:
                      # object already exists, previous submission was successful
                      # even though it errored
                      break
                  raise
              except urllib3.exceptions.HTTPError:
                  attempts -= 1
                  if attempts == 0:
                      raise
                  self.logger.warning(
                      "Error submitting job %s, retrying...", job_name, exc_info=True
                  )
                  time.sleep(1)
  
          self.logger.debug("Job {} created".format(job_name))
  
          return "Job {}".format(job_name)
  
      def generate_job_spec(self, flow_run: GraphQLResult) -> dict:
          """Generate a k8s job spec for a flow run
  
          Args:
              - flow_run (GraphQLResult): A flow run object
  
          Returns:
              - dict: a dictionary representation of a k8s job for flow execution
          """
  <<<<<<< HEAD
          run_config = self._get_run_config(flow_run, KubernetesRun)
          assert run_config is None or isinstance(run_config, KubernetesRun)  # mypy
          if run_config is not None:
              return self.generate_job_spec_from_run_config(flow_run, run_config)
  =======
          if getattr(flow_run.flow, "run_config", None) is not None:
              return self.generate_job_spec_from_run_config(flow_run)
  >>>>>>> prefect clone
          else:
              return self.generate_job_spec_from_environment(flow_run)
  
      def generate_job_spec_from_environment(
          self, flow_run: GraphQLResult, image: str = None
      ) -> dict:
          """
          Populate a k8s job spec. This spec defines a k8s job that handles
          executing a flow. This method runs each time the agent receives
          a flow to run.
  
          That job spec can optionally be customized by setting the
          following environment variables on the agent.
  
          - `NAMESPACE`: the k8s namespace the job will run in, defaults to `"default"`
          - `JOB_MEM_REQUEST`: memory requested, for example, `256Mi` for 256 MB. If this
                  environment variable is not set, the cluster's defaults will be used.
          - `JOB_MEM_LIMIT`: memory limit, for example, `512Mi` For 512 MB. If this
                  environment variable is not set, the cluster's defaults will be used.
          - `JOB_CPU_REQUEST`: CPU requested, defaults to `"100m"`
          - `JOB_CPU_LIMIT`: CPU limit, defaults to `"100m"`
          - `IMAGE_PULL_POLICY`: policy for pulling images. Defaults to `"IfNotPresent"`.
          - `IMAGE_PULL_SECRETS`: name of an existing k8s secret that can be used to pull
                  images. This is necessary if your flow uses an image that is in a non-public
                  container registry, such as Amazon ECR, or in a public registry that requires
                  authentication to avoid hitting rate limits. To specify multiple image pull
                  secrets, provide a comma-delimited string with no spaces, like
                  `"some-secret,other-secret"`.
          - `SERVICE_ACCOUNT_NAME`: name of a service account to run the job as.
                  By default, none is specified.
          - `YAML_TEMPLATE`: a path to where the YAML template should be loaded from. defaults
                  to the embedded `job_spec.yaml`.
  
          Args:
              - flow_run (GraphQLResult): A flow run object
              - image (str, optional): The full name of an image to use for the job
  
          Returns:
              - dict: a dictionary representation of a k8s job for flow execution
          """
          identifier = str(uuid.uuid4())[:8]
          yaml_path = os.getenv(
              "YAML_TEMPLATE", os.path.join(os.path.dirname(__file__), "job_spec.yaml")
          )
          with open(yaml_path, "r") as job_file:
              job = yaml.safe_load(job_file)
  
          job_name = "prefect-job-{}".format(identifier)
  
          # Populate job metadata for identification
          k8s_labels = {
              "prefect.io/identifier": identifier,
              "prefect.io/flow_run_id": flow_run.id,  # type: ignore
              "prefect.io/flow_id": flow_run.flow.id,  # type: ignore
          }
          job["metadata"]["name"] = job_name
          job["metadata"]["labels"].update(**k8s_labels)
          job["spec"]["template"]["metadata"]["labels"].update(**k8s_labels)
  
          # Use provided image for job
          if image is None:
              image = get_flow_image(flow_run=flow_run)
          job["spec"]["template"]["spec"]["containers"][0]["image"] = image
  
          self.logger.debug("Using image {} for job".format(image))
  
          # Datermine flow run command
          job["spec"]["template"]["spec"]["containers"][0]["args"] = [
              get_flow_run_command(flow_run)
          ]
  
          # Populate environment variables for flow run execution
          env = job["spec"]["template"]["spec"]["containers"][0]["env"]
  
          env[0]["value"] = config.cloud.api or "https://api.prefect.io"
          env[1]["value"] = config.cloud.agent.auth_token
          env[2]["value"] = flow_run.id  # type: ignore
          env[3]["value"] = flow_run.flow.id  # type: ignore
  <<<<<<< HEAD
          env[4]["value"] = self.namespace
  =======
          env[4]["value"] = os.getenv("NAMESPACE", "default")
  >>>>>>> prefect clone
          env[5]["value"] = str(self.labels)
          env[6]["value"] = str(self.log_to_cloud).lower()
          env[7]["value"] = config.logging.level
  
          # append all user provided values
          for key, value in self.env_vars.items():
              env.append(dict(name=key, value=value))
  
          # Use image pull secrets if provided
  <<<<<<< HEAD
          if self.image_pull_secrets:
              for idx, secret_name in enumerate(self.image_pull_secrets):
  =======
          image_pull_secrets = os.getenv("IMAGE_PULL_SECRETS")
          if image_pull_secrets:
              secrets = image_pull_secrets.split(",")
              for idx, secret_name in enumerate(secrets):
  >>>>>>> prefect clone
                  # this check preserves behavior from previous releases,
                  # where prefect would only overwrite the first entry in
                  # imagePullSecrets
                  if idx == 0:
                      job["spec"]["template"]["spec"]["imagePullSecrets"][0] = {
                          "name": secret_name
                      }
                  else:
                      job["spec"]["template"]["spec"]["imagePullSecrets"].append(
                          {"name": secret_name}
                      )
          else:
              del job["spec"]["template"]["spec"]["imagePullSecrets"]
  
          # Set resource requirements if provided
          resources = job["spec"]["template"]["spec"]["containers"][0]["resources"]
          if os.getenv("JOB_MEM_REQUEST"):
              resources["requests"]["memory"] = os.getenv("JOB_MEM_REQUEST")
          if os.getenv("JOB_MEM_LIMIT"):
              resources["limits"]["memory"] = os.getenv("JOB_MEM_LIMIT")
          if os.getenv("JOB_CPU_REQUEST"):
              resources["requests"]["cpu"] = os.getenv("JOB_CPU_REQUEST")
          if os.getenv("JOB_CPU_LIMIT"):
              resources["limits"]["cpu"] = os.getenv("JOB_CPU_LIMIT")
          if self.volume_mounts:
              job["spec"]["template"]["spec"]["containers"][0][
                  "volumeMounts"
              ] = self.volume_mounts
          else:
              del job["spec"]["template"]["spec"]["containers"][0]["volumeMounts"]
          if self.volumes:
              job["spec"]["template"]["spec"]["volumes"] = self.volumes
          else:
              del job["spec"]["template"]["spec"]["volumes"]
          if os.getenv("IMAGE_PULL_POLICY"):
              job["spec"]["template"]["spec"]["containers"][0][
                  "imagePullPolicy"
              ] = os.getenv("IMAGE_PULL_POLICY")
  <<<<<<< HEAD
          if self.service_account_name:
              job["spec"]["template"]["spec"][
                  "serviceAccountName"
              ] = self.service_account_name
  
          return job
  
      def generate_job_spec_from_run_config(
          self, flow_run: GraphQLResult, run_config: KubernetesRun
      ) -> dict:
  =======
          if os.getenv("SERVICE_ACCOUNT_NAME"):
              job["spec"]["template"]["spec"]["serviceAccountName"] = os.getenv(
                  "SERVICE_ACCOUNT_NAME"
              )
  
          return job
  
      def generate_job_spec_from_run_config(self, flow_run: GraphQLResult) -> dict:
  >>>>>>> prefect clone
          """Generate a k8s job spec for a flow run.
  
          Args:
              - flow_run (GraphQLResult): A flow run object
  <<<<<<< HEAD
              - run_config (KubernetesRun): The flow run's run_config
  =======
  >>>>>>> prefect clone
  
          Returns:
              - dict: a dictionary representation of a k8s job for flow execution
          """
  <<<<<<< HEAD
  =======
          run_config = RunConfigSchema().load(flow_run.flow.run_config)
  
  >>>>>>> prefect clone
          if run_config.job_template:
              job = run_config.job_template
          else:
              job_template_path = run_config.job_template_path or self.job_template_path
              self.logger.debug("Loading job template from %r", job_template_path)
              template_bytes = read_bytes_from_path(job_template_path)
              job = yaml.safe_load(template_bytes)
  
          identifier = uuid.uuid4().hex[:8]
  
          job_name = f"prefect-job-{identifier}"
  
          # Populate job metadata for identification
          k8s_labels = {
              "prefect.io/identifier": identifier,
              "prefect.io/flow_run_id": flow_run.id,  # type: ignore
              "prefect.io/flow_id": flow_run.flow.id,  # type: ignore
          }
          _get_or_create(job, "metadata.labels")
          _get_or_create(job, "spec.template.metadata.labels")
          job["metadata"]["name"] = job_name
          job["metadata"]["labels"].update(**k8s_labels)
          job["spec"]["template"]["metadata"]["labels"].update(**k8s_labels)
  <<<<<<< HEAD
          pod_spec = job["spec"]["template"]["spec"]
  
          # Configure `service_account_name` if specified
          if run_config.service_account_name is not None:
              # On run-config, always override
              service_account_name = (
                  run_config.service_account_name
              )  # type: Optional[str]
          elif "serviceAccountName" in pod_spec and (
              run_config.job_template or run_config.job_template_path
          ):
              # On run-config job-template, no override
              service_account_name = None
          else:
              # Use agent value, if provided
              service_account_name = self.service_account_name
          if service_account_name is not None:
              pod_spec["serviceAccountName"] = service_account_name
  
          # Configure `image_pull_secrets` if specified
          if run_config.image_pull_secrets is not None:
              # On run-config, always override
              image_pull_secrets = (
                  run_config.image_pull_secrets
              )  # type: Optional[Iterable[str]]
          elif "imagePullSecrets" in pod_spec and (
              run_config.job_template or run_config.job_template_path
          ):
              # On run-config job template, no override
              image_pull_secrets = None
          else:
              # Use agent, if provided
              image_pull_secrets = self.image_pull_secrets
          if image_pull_secrets is not None:
              pod_spec["imagePullSecrets"] = [{"name": s} for s in image_pull_secrets]
  
          # Default restartPolicy to Never
          _get_or_create(job, "spec.template.spec.restartPolicy", "Never")
  =======
  >>>>>>> prefect clone
  
          # Get the first container, which is used for the prefect job
          containers = _get_or_create(job, "spec.template.spec.containers", [])
          if not containers:
              containers.append({})
          container = containers[0]
  
          # Set container image
  <<<<<<< HEAD
          container["image"] = image = get_flow_image(flow_run)
  
          # Set flow run command
          container["args"] = get_flow_run_command(flow_run).split()
  =======
          container["image"] = get_flow_image(flow_run)
  
          # Set flow run command
          container["args"] = [get_flow_run_command(flow_run)]
  >>>>>>> prefect clone
  
          # Populate environment variables from the following sources,
          # with precedence:
          # - Values required for flow execution, hardcoded below
          # - Values set on the KubernetesRun object
          # - Values set using the `--env` CLI flag on the agent
          # - Values in the job template
          env = self.env_vars.copy()
          if run_config.env:
              env.update(run_config.env)
          env.update(
              {
  <<<<<<< HEAD
                  "PREFECT__CLOUD__AGENT__LABELS": str(self.labels),
  =======
  >>>>>>> prefect clone
                  "PREFECT__CLOUD__API": config.cloud.api,
                  "PREFECT__CLOUD__AUTH_TOKEN": config.cloud.agent.auth_token,
                  "PREFECT__CLOUD__USE_LOCAL_SECRETS": "false",
                  "PREFECT__CONTEXT__FLOW_RUN_ID": flow_run.id,
                  "PREFECT__CONTEXT__FLOW_ID": flow_run.flow.id,
  <<<<<<< HEAD
                  "PREFECT__CONTEXT__IMAGE": image,
  =======
  >>>>>>> prefect clone
                  "PREFECT__LOGGING__LOG_TO_CLOUD": str(self.log_to_cloud).lower(),
                  "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudFlowRunner",
                  "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS": "prefect.engine.cloud.CloudTaskRunner",
              }
          )
          container_env = [{"name": k, "value": v} for k, v in env.items()]
          for entry in container.get("env", []):
              if entry["name"] not in env:
                  container_env.append(entry)
          container["env"] = container_env
  
          # Set resource requirements if provided
          _get_or_create(container, "resources.requests")
          _get_or_create(container, "resources.limits")
          resources = container["resources"]
          if run_config.memory_request:
              resources["requests"]["memory"] = run_config.memory_request
          if run_config.memory_limit:
              resources["limits"]["memory"] = run_config.memory_limit
          if run_config.cpu_request:
              resources["requests"]["cpu"] = run_config.cpu_request
          if run_config.cpu_limit:
              resources["limits"]["cpu"] = run_config.cpu_limit
  
          return job
  
      @staticmethod
      def generate_deployment_yaml(
          token: str = None,
          api: str = None,
          namespace: str = None,
          image_pull_secrets: str = None,
  <<<<<<< HEAD
  =======
          resource_manager_enabled: bool = False,
  >>>>>>> prefect clone
          rbac: bool = False,
          latest: bool = False,
          mem_request: str = None,
          mem_limit: str = None,
          cpu_request: str = None,
          cpu_limit: str = None,
          image_pull_policy: str = None,
          service_account_name: str = None,
          labels: Iterable[str] = None,
          env_vars: dict = None,
          backend: str = None,
      ) -> str:
          """
          Generate and output an installable YAML spec for the agent.
  
          Args:
              - token (str, optional): A `RUNNER` token to give the agent
              - api (str, optional): A URL pointing to the Prefect API. Defaults to
                  `https://api.prefect.io`
              - namespace (str, optional): The namespace to create Prefect jobs in. Defaults
                  to `default`
              - image_pull_secrets (str, optional): The name of an image pull secret to use
                  for Prefect jobs
  <<<<<<< HEAD
  =======
              - resource_manager_enabled (bool, optional): Whether to include the resource
                  manager as part of the YAML. Defaults to `False`
  >>>>>>> prefect clone
              - rbac (bool, optional): Whether to include default RBAC configuration as
                  part of the YAML. Defaults to `False`
              - latest (bool, optional): Whether to use the `latest` Prefect image.
                  Defaults to `False`
              - mem_request (str, optional): Requested memory for Prefect init job.
              - mem_limit (str, optional): Limit memory for Prefect init job.
              - cpu_request (str, optional): Requested CPU for Prefect init job.
              - cpu_limit (str, optional): Limit CPU for Prefect init job.
              - image_pull_policy (str, optional): imagePullPolicy to use for Prefect init job.
                  Job defaults to `IfNotPresent`.
              - service_account_name (str, optional): Name of a service account to use for
                  Prefect init job. Job defaults to using `default` service account.
              - labels (List[str], optional): a list of labels, which are arbitrary string
                  identifiers used by Prefect Agents when polling for work
              - env_vars (dict, optional): additional environment variables to attach to all
                  jobs created by this agent and to set in the agent's own environment
              - backend (str, optional): toggle which backend to use for this agent.
                  Defaults to backend currently set in config.
  
          Returns:
              - str: A string representation of the generated YAML
          """
  
          # Use defaults if not provided
          token = token or ""
          api = api or "https://api.prefect.io"
          namespace = namespace or "default"
          labels = labels or []
          mem_request = mem_request or ""
          mem_limit = mem_limit or ""
          cpu_request = cpu_request or ""
          cpu_limit = cpu_limit or ""
          image_pull_policy = image_pull_policy or ""
          service_account_name = service_account_name or ""
          backend = backend or config.backend
  
          version = prefect.__version__.split("+")
          image_version = (
              "latest" if len(version) > 1 or latest else (version[0] + "-python3.6")
          )
  
          with open(
              os.path.join(os.path.dirname(__file__), "deployment.yaml"), "r"
          ) as deployment_file:
              deployment = yaml.safe_load(deployment_file)
  
          agent_env = deployment["spec"]["template"]["spec"]["containers"][0]["env"]
  
          # Populate env vars
          agent_env[0]["value"] = token
          agent_env[1]["value"] = api
          agent_env[2]["value"] = namespace
          agent_env[3]["value"] = image_pull_secrets or ""
          agent_env[4]["value"] = str(labels)
          agent_env[11]["value"] = backend
  
          # Populate job resource env vars
          agent_env[5]["value"] = mem_request
          agent_env[6]["value"] = mem_limit
          agent_env[7]["value"] = cpu_request
          agent_env[8]["value"] = cpu_limit
          agent_env[9]["value"] = image_pull_policy
          agent_env[10]["value"] = service_account_name
  
          if env_vars:
              agent_env.append(
                  {
                      "name": "PREFECT__CLOUD__AGENT__ENV_VARS",
                      "value": json.dumps(env_vars),
                  }
              )
  
          # Use local prefect version for image
          deployment["spec"]["template"]["spec"]["containers"][0][
              "image"
          ] = "prefecthq/prefect:{}".format(image_version)
  
  <<<<<<< HEAD
  =======
          # Populate resource manager if requested
          if resource_manager_enabled:
              resource_manager_env = deployment["spec"]["template"]["spec"]["containers"][
                  1
              ]["env"]
  
              resource_manager_env[0]["value"] = token
              resource_manager_env[1]["value"] = api
              resource_manager_env[3]["value"] = namespace
  
              # Use local prefect version for image
              deployment["spec"]["template"]["spec"]["containers"][1][
                  "image"
              ] = "prefecthq/prefect:{}".format(image_version)
          else:
              del deployment["spec"]["template"]["spec"]["containers"][1]
  
  >>>>>>> prefect clone
          # Populate image pull secrets if provided
          if image_pull_secrets:
              agent_env = deployment["spec"]["template"]["spec"]["imagePullSecrets"][0][
                  "name"
              ] = image_pull_secrets
          else:
              del deployment["spec"]["template"]["spec"]["imagePullSecrets"]
  
          # Load RBAC if specified
          rbac_yaml = []
          if rbac:
              with open(
                  os.path.join(os.path.dirname(__file__), "rbac.yaml"), "r"
              ) as rbac_file:
                  rbac_generator = yaml.safe_load_all(rbac_file)
  
                  for document in rbac_generator:
                      document["metadata"]["namespace"] = namespace
                      rbac_yaml.append(document)
  
          output_yaml = [deployment]
          output_yaml.extend(rbac_yaml)
          return yaml.safe_dump_all(output_yaml, explicit_start=True)
  
  
  if __name__ == "__main__":
      KubernetesAgent().start()
  