Repository: PrefectHQ/prefect

Issue:
How should environments be loaded in Cloud?
Currently, environments are loaded via marshmallow deserialization here: https://github.com/PrefectHQ/prefect/blob/master/src/prefect/cli/execute.py#L66-L67

This works great, except that it prevents users from implementing custom execution environments.  Because this CLI typically runs inside a docker container containing the Flow bytecode, we could load that, etc. but then we're taking over the responsibilities of `environment.execute()`, which IMO means we have a leaky abstraction somewhere.

Our ideal solution to this problem should support:
- custom user environments
- the (future) ability to specify environments on a per-flow-run basis
- labels on custom environments should be readable on the Cloud-side through the marshmallow deserialization process


Discussion:
What do you mean when you say _we're taking over the responsibilities of `environment.execute()`_?
`environment.execute` currently accepts a flow location string (which is typically the filesystem location of the flow in the docker container), and loads the flow, etc.

If the CLI starts performing that loading from disk, they now share mixed responsibilities - so we'd either have to break the environment interface or redundantly reload the flow.
As discussed with @joshmeek we could:
- load the storage from the Cloud serialization protocol
- call `get_flow` to retrieve the actual flow object
- use `flow.environment` directly, which would allow for user-written environments

The interface becomes a little redundant (loading the flow twice...) but we could work around that by introducing a new keyword argument to `environment.execute` which accepts an _actual_ flow object and bypasses the load.  This would allow for environments which need to load the flow themselves, or environments which can accept flow objects directly.

Resolving patch:
diff --git a/src/prefect/cli/execute.py b/src/prefect/cli/execute.py
--- a/src/prefect/cli/execute.py
+++ b/src/prefect/cli/execute.py
@@ -63,8 +63,8 @@ def cloud_flow():
         storage_schema = prefect.serialization.storage.StorageSchema()
         storage = storage_schema.load(flow_data.storage)
 
-        environment_schema = prefect.serialization.environment.EnvironmentSchema()
-        environment = environment_schema.load(flow_data.environment)
+        flow = storage.get_flow(storage.flows[flow_data.name])
+        environment = flow.environment
 
         environment.setup(storage=storage)
         environment.execute(
diff --git a/src/prefect/environments/execution/dask/k8s.py b/src/prefect/environments/execution/dask/k8s.py
--- a/src/prefect/environments/execution/dask/k8s.py
+++ b/src/prefect/environments/execution/dask/k8s.py
@@ -1,21 +1,16 @@
 import base64
 import json
-import sys
-import time
 import uuid
 from os import path
 from typing import Any, List
 
 import cloudpickle
-import docker
 import yaml
-from slugify import slugify
 
 import prefect
 from prefect.client import Secret
 from prefect.environments.execution import Environment
 from prefect.environments.storage import Docker
-from prefect.utilities import logging
 
 
 class DaskKubernetesEnvironment(Environment):
@@ -30,6 +25,18 @@ class DaskKubernetesEnvironment(Environment):
     function does have the requirement in that it needs to have an `identifier_label`
     set with a UUID so resources can be cleaned up independently of other deployments.
 
+    It is possible to provide a custom scheduler and worker spec YAML files through the `scheduler_spec_file` and
+    `worker_spec_file` arguments. These specs (if provided) will be used in place of the defaults. Your spec files
+    should be modeled after the job.yaml and worker_pod.yaml found at
+    https://github.com/PrefectHQ/prefect/tree/master/src/prefect/environments/execution/dask
+    The main aspects to be aware of are the `command` and `args` on the container. These environment variables are
+    required for cloud do not need to be included because they are instead automatically added and populated during execution:
+    PREFECT__CLOUD__GRAPHQL, PREFECT__CLOUD__AUTH_TOKEN, PREFECT__CONTEXT__FLOW_RUN_ID,
+    PREFECT__CONTEXT__NAMESPACE, PREFECT__CONTEXT__IMAGE, PREFECT__CONTEXT__FLOW_FILE_PATH,
+    PREFECT__CLOUD__USE_LOCAL_SECRETS, PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS,
+    PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS, PREFECT__ENGINE__EXECUTOR__DEFAULT_CLASS,
+    PREFECT__LOGGING__LOG_TO_CLOUD
+
     Args:
         - min_workers (int, optional): the minimum allowed number of Dask worker pods; defaults to 1
         - max_workers (int, optional): the maximum allowed number of Dask worker pods; defaults to 1
@@ -41,6 +48,8 @@ class DaskKubernetesEnvironment(Environment):
             `"docker-username"`, `"docker-password"`, and `"docker-email"`.
         - labels (List[str], optional): a list of labels, which are arbitrary string identifiers used by Prefect
             Agents when polling for work
+        - scheduler_spec_file (str, optional): Path to a scheduler spec YAML file
+        - worker_spec_file (str, optional): Path to a worker spec YAML file
     """
 
     def __init__(
@@ -50,6 +59,8 @@ def __init__(
         private_registry: bool = False,
         docker_secret: str = None,
         labels: List[str] = None,
+        scheduler_spec_file: str = None,
+        worker_spec_file: str = None,
     ) -> None:
         self.min_workers = min_workers
         self.max_workers = max_workers
@@ -59,6 +70,12 @@ def __init__(
             self.docker_secret = docker_secret or "DOCKER_REGISTRY_CREDENTIALS"
         else:
             self.docker_secret = None  # type: ignore
+        self.scheduler_spec_file = scheduler_spec_file
+        self.worker_spec_file = worker_spec_file
+
+        # Load specs from file if path given, store on object
+        self._scheduler_spec, self._worker_spec = self._load_specs_from_file()
+
         super().__init__(labels=labels)
 
     def setup(self, storage: "Docker") -> None:  # type: ignore
@@ -179,20 +196,26 @@ def create_flow_run_job(self, docker_name: str, flow_file_path: str) -> None:
 
         batch_client = client.BatchV1Api()
 
-        with open(path.join(path.dirname(__file__), "job.yaml")) as job_file:
-            job = yaml.safe_load(job_file)
-            job = self._populate_job_yaml(
+        if self._scheduler_spec:
+            job = self._scheduler_spec
+            job = self._populate_scheduler_spec_yaml(
                 yaml_obj=job, docker_name=docker_name, flow_file_path=flow_file_path
             )
-
-            # Create Job
-            try:
-                batch_client.create_namespaced_job(
-                    namespace=prefect.context.get("namespace"), body=job
+        else:
+            with open(path.join(path.dirname(__file__), "job.yaml")) as job_file:
+                job = yaml.safe_load(job_file)
+                job = self._populate_job_yaml(
+                    yaml_obj=job, docker_name=docker_name, flow_file_path=flow_file_path
                 )
-            except Exception as exc:
-                self.logger.critical("Failed to create Kubernetes job: {}".format(exc))
-                raise exc
+
+        # Create Job
+        try:
+            batch_client.create_namespaced_job(
+                namespace=prefect.context.get("namespace"), body=job
+            )
+        except Exception as exc:
+            self.logger.critical("Failed to create Kubernetes job: {}".format(exc))
+            raise exc
 
     def run_flow(self) -> None:
         """
@@ -203,36 +226,42 @@ def run_flow(self) -> None:
             from prefect.engine.executors import DaskExecutor
             from dask_kubernetes import KubeCluster
 
-            with open(path.join(path.dirname(__file__), "worker_pod.yaml")) as pod_file:
-                worker_pod = yaml.safe_load(pod_file)
-                worker_pod = self._populate_worker_pod_yaml(yaml_obj=worker_pod)
-
-                cluster = KubeCluster.from_dict(
-                    worker_pod, namespace=prefect.context.get("namespace")
-                )
-                cluster.adapt(minimum=self.min_workers, maximum=self.max_workers)
-
-                # Load serialized flow from file and run it with a DaskExecutor
+            if self._worker_spec:
+                worker_pod = self._worker_spec
+                worker_pod = self._populate_worker_spec_yaml(yaml_obj=worker_pod)
+            else:
                 with open(
-                    prefect.context.get(
-                        "flow_file_path", "/root/.prefect/flow_env.prefect"
-                    ),
-                    "rb",
-                ) as f:
-                    flow = cloudpickle.load(f)
-
-                    executor = DaskExecutor(address=cluster.scheduler_address)
-                    runner_cls = get_default_flow_runner_class()
-                    runner_cls(flow=flow).run(executor=executor)
+                    path.join(path.dirname(__file__), "worker_pod.yaml")
+                ) as pod_file:
+                    worker_pod = yaml.safe_load(pod_file)
+                    worker_pod = self._populate_worker_pod_yaml(yaml_obj=worker_pod)
+
+            cluster = KubeCluster.from_dict(
+                worker_pod, namespace=prefect.context.get("namespace")
+            )
+            cluster.adapt(minimum=self.min_workers, maximum=self.max_workers)
+
+            # Load serialized flow from file and run it with a DaskExecutor
+            with open(
+                prefect.context.get(
+                    "flow_file_path", "/root/.prefect/flow_env.prefect"
+                ),
+                "rb",
+            ) as f:
+                flow = cloudpickle.load(f)
+
+                executor = DaskExecutor(address=cluster.scheduler_address)
+                runner_cls = get_default_flow_runner_class()
+                runner_cls(flow=flow).run(executor=executor)
         except Exception as exc:
             self.logger.exception(
                 "Unexpected error raised during flow run: {}".format(exc)
             )
             raise exc
 
-    ########################
-    # YAML Spec Manipulation
-    ########################
+    ################################
+    # Default YAML Spec Manipulation
+    ################################
 
     def _populate_job_yaml(
         self, yaml_obj: dict, docker_name: str, flow_file_path: str
@@ -315,3 +344,142 @@ def _populate_worker_pod_yaml(self, yaml_obj: dict) -> dict:
         )
 
         return yaml_obj
+
+    ###############################
+    # Custom YAML Spec Manipulation
+    ###############################
+
+    def _populate_scheduler_spec_yaml(
+        self, yaml_obj: dict, docker_name: str, flow_file_path: str
+    ) -> dict:
+        """
+        Populate the custom execution job yaml object used in this environment with the proper values
+
+        Args:
+            - yaml_obj (dict): A dictionary representing the parsed yaml
+            - docker_name (str): the full path to the docker image
+            - flow_file_path (str): the location of the flow within the docker container
+
+        Returns:
+            - dict: a dictionary with the yaml values replaced
+        """
+        flow_run_id = prefect.context.get("flow_run_id", "unknown")
+
+        yaml_obj["metadata"]["labels"]["identifier"] = self.identifier_label
+        yaml_obj["metadata"]["labels"]["flow_run_id"] = flow_run_id
+        yaml_obj["spec"]["template"]["metadata"]["labels"][
+            "identifier"
+        ] = self.identifier_label
+
+        # Required Cloud environment variables
+        env_values = [
+            {"name": "PREFECT__CLOUD__GRAPHQL", "value": prefect.config.cloud.graphql},
+            {
+                "name": "PREFECT__CLOUD__AUTH_TOKEN",
+                "value": prefect.config.cloud.auth_token,
+            },
+            {"name": "PREFECT__CONTEXT__FLOW_RUN_ID", "value": flow_run_id},
+            {
+                "name": "PREFECT__CONTEXT__NAMESPACE",
+                "value": prefect.context.get("namespace", ""),
+            },
+            {"name": "PREFECT__CONTEXT__IMAGE", "value": docker_name},
+            {"name": "PREFECT__CONTEXT__FLOW_FILE_PATH", "value": flow_file_path},
+            {"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false"},
+            {
+                "name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
+                "value": "prefect.engine.cloud.CloudFlowRunner",
+            },
+            {
+                "name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
+                "value": "prefect.engine.cloud.CloudTaskRunner",
+            },
+            {
+                "name": "PREFECT__ENGINE__EXECUTOR__DEFAULT_CLASS",
+                "value": "prefect.engine.executors.DaskExecutor",
+            },
+            {"name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": "true"},
+        ]
+
+        # set environment variables
+        env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"]
+        env.extend(env_values)
+
+        # set image
+        yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"] = docker_name
+
+        return yaml_obj
+
+    def _populate_worker_spec_yaml(self, yaml_obj: dict) -> dict:
+        """
+        Populate the custom worker pod yaml object used in this environment with the proper values.
+
+        Args:
+            - yaml_obj (dict): A dictionary representing the parsed yaml
+
+        Returns:
+            - dict: a dictionary with the yaml values replaced
+        """
+        # set identifier labels
+        yaml_obj["metadata"]["labels"]["identifier"] = self.identifier_label
+        yaml_obj["metadata"]["labels"]["flow_run_id"] = prefect.context.get(
+            "flow_run_id", "unknown"
+        )
+
+        # Required Cloud environment variables
+        env_values = [
+            {"name": "PREFECT__CLOUD__GRAPHQL", "value": prefect.config.cloud.graphql},
+            {
+                "name": "PREFECT__CLOUD__AUTH_TOKEN",
+                "value": prefect.config.cloud.auth_token,
+            },
+            {
+                "name": "PREFECT__CONTEXT__FLOW_RUN_ID",
+                "value": prefect.context.get("flow_run_id", ""),
+            },
+            {"name": "PREFECT__CLOUD__USE_LOCAL_SECRETS", "value": "false"},
+            {
+                "name": "PREFECT__ENGINE__FLOW_RUNNER__DEFAULT_CLASS",
+                "value": "prefect.engine.cloud.CloudFlowRunner",
+            },
+            {
+                "name": "PREFECT__ENGINE__TASK_RUNNER__DEFAULT_CLASS",
+                "value": "prefect.engine.cloud.CloudTaskRunner",
+            },
+            {
+                "name": "PREFECT__ENGINE__EXECUTOR__DEFAULT_CLASS",
+                "value": "prefect.engine.executors.DaskExecutor",
+            },
+            {"name": "PREFECT__LOGGING__LOG_TO_CLOUD", "value": "true"},
+        ]
+
+        # set environment variables
+        env = yaml_obj["spec"]["containers"][0]["env"]
+        env.extend(env_values)
+
+        # set image
+        yaml_obj["spec"]["containers"][0]["image"] = prefect.context.get(
+            "image", "daskdev/dask:latest"
+        )
+
+        return yaml_obj
+
+    def _load_specs_from_file(self) -> tuple:
+        """
+        Load scheduler and worker spec from provided file paths
+
+        Returns:
+            - tuple: scheduler spec dictionary, worker spec dictionary
+        """
+        scheduler = None
+        worker = None
+
+        if self.scheduler_spec_file:
+            with open(self.scheduler_spec_file) as scheduler_spec_file:
+                scheduler = yaml.safe_load(scheduler_spec_file)
+
+        if self.worker_spec_file:
+            with open(self.worker_spec_file) as worker_spec_file:
+                worker = yaml.safe_load(worker_spec_file)
+
+        return scheduler, worker
diff --git a/src/prefect/serialization/environment.py b/src/prefect/serialization/environment.py
--- a/src/prefect/serialization/environment.py
+++ b/src/prefect/serialization/environment.py
@@ -33,6 +33,10 @@ class Meta:
     min_workers = fields.Int()
     max_workers = fields.Int()
 
+    # Serialize original spec file locations, not the spec itself
+    scheduler_spec_file = fields.String(allow_none=True)
+    worker_spec_file = fields.String(allow_none=True)
+
 
 class RemoteEnvironmentSchema(ObjectSchema):
     class Meta:


Test patch:
diff --git a/tests/environments/execution/test_dask_k8s_environment.py b/tests/environments/execution/test_dask_k8s_environment.py
--- a/tests/environments/execution/test_dask_k8s_environment.py
+++ b/tests/environments/execution/test_dask_k8s_environment.py
@@ -259,3 +259,112 @@ def test_populate_worker_pod_yaml_with_private_registry():
             yaml_obj = environment._populate_worker_pod_yaml(yaml_obj=pod)
 
     yaml_obj["spec"]["imagePullSecrets"][0] == dict(name="foo-man-docker")
+
+
+def test_initialize_environment_with_spec_populates(monkeypatch):
+
+    with tempfile.TemporaryDirectory() as directory:
+
+        with open(os.path.join(directory, "scheduler.yaml"), "w+") as file:
+            file.write("scheduler")
+        with open(os.path.join(directory, "worker.yaml"), "w+") as file:
+            file.write("worker")
+
+        environment = DaskKubernetesEnvironment(
+            scheduler_spec_file=os.path.join(directory, "scheduler.yaml"),
+            worker_spec_file=os.path.join(directory, "worker.yaml"),
+        )
+
+        assert environment._scheduler_spec == "scheduler"
+        assert environment._worker_spec == "worker"
+
+
+def test_populate_custom_worker_spec_yaml():
+    environment = DaskKubernetesEnvironment()
+
+    file_path = os.path.dirname(prefect.environments.execution.dask.k8s.__file__)
+
+    with open(path.join(file_path, "worker_pod.yaml")) as pod_file:
+        pod = yaml.safe_load(pod_file)
+        pod["spec"]["containers"][0]["env"] = []
+
+    with set_temporary_config(
+        {"cloud.graphql": "gql_test", "cloud.auth_token": "auth_test"}
+    ):
+        with prefect.context(flow_run_id="id_test", image="my_image"):
+            yaml_obj = environment._populate_worker_spec_yaml(yaml_obj=pod)
+
+    assert yaml_obj["metadata"]["labels"]["identifier"] == environment.identifier_label
+    assert yaml_obj["metadata"]["labels"]["flow_run_id"] == "id_test"
+
+    env = yaml_obj["spec"]["containers"][0]["env"]
+
+    assert env[0]["value"] == "gql_test"
+    assert env[1]["value"] == "auth_test"
+    assert env[2]["value"] == "id_test"
+    assert env[3]["value"] == "false"
+    assert env[4]["value"] == "prefect.engine.cloud.CloudFlowRunner"
+    assert env[5]["value"] == "prefect.engine.cloud.CloudTaskRunner"
+    assert env[6]["value"] == "prefect.engine.executors.DaskExecutor"
+    assert env[7]["value"] == "true"
+
+    assert yaml_obj["spec"]["containers"][0]["image"] == "my_image"
+
+
+def test_populate_custom_scheduler_spec_yaml():
+    environment = DaskKubernetesEnvironment()
+
+    file_path = os.path.dirname(prefect.environments.execution.dask.k8s.__file__)
+
+    with open(path.join(file_path, "job.yaml")) as job_file:
+        job = yaml.safe_load(job_file)
+        job["spec"]["template"]["spec"]["containers"][0]["env"] = []
+
+    with set_temporary_config(
+        {"cloud.graphql": "gql_test", "cloud.auth_token": "auth_test"}
+    ):
+        with prefect.context(flow_run_id="id_test", namespace="namespace_test"):
+            yaml_obj = environment._populate_scheduler_spec_yaml(
+                yaml_obj=job, docker_name="test1/test2:test3", flow_file_path="test4"
+            )
+
+    env = yaml_obj["spec"]["template"]["spec"]["containers"][0]["env"]
+
+    assert env[0]["value"] == "gql_test"
+    assert env[1]["value"] == "auth_test"
+    assert env[2]["value"] == "id_test"
+    assert env[3]["value"] == "namespace_test"
+    assert env[4]["value"] == "test1/test2:test3"
+    assert env[5]["value"] == "test4"
+    assert env[6]["value"] == "false"
+    assert env[7]["value"] == "prefect.engine.cloud.CloudFlowRunner"
+    assert env[8]["value"] == "prefect.engine.cloud.CloudTaskRunner"
+    assert env[9]["value"] == "prefect.engine.executors.DaskExecutor"
+    assert env[10]["value"] == "true"
+
+    assert (
+        yaml_obj["spec"]["template"]["spec"]["containers"][0]["image"]
+        == "test1/test2:test3"
+    )
+
+
+def test_roundtrip_cloudpickle():
+    with tempfile.TemporaryDirectory() as directory:
+
+        with open(os.path.join(directory, "scheduler.yaml"), "w+") as file:
+            file.write("scheduler")
+        with open(os.path.join(directory, "worker.yaml"), "w+") as file:
+            file.write("worker")
+
+        environment = DaskKubernetesEnvironment(
+            scheduler_spec_file=os.path.join(directory, "scheduler.yaml"),
+            worker_spec_file=os.path.join(directory, "worker.yaml"),
+        )
+
+        assert environment._scheduler_spec == "scheduler"
+        assert environment._worker_spec == "worker"
+
+        new = cloudpickle.loads(cloudpickle.dumps(environment))
+        assert isinstance(new, DaskKubernetesEnvironment)
+        assert new._scheduler_spec == "scheduler"
+        assert new._worker_spec == "worker"
diff --git a/tests/serialization/test_environments.py b/tests/serialization/test_environments.py
--- a/tests/serialization/test_environments.py
+++ b/tests/serialization/test_environments.py
@@ -38,6 +38,8 @@ def test_serialize_dask_environment():
     assert serialized["min_workers"] == 1
     assert serialized["max_workers"] == 2
     assert serialized["labels"] == []
+    assert serialized["scheduler_spec_file"] is None
+    assert serialized["worker_spec_file"] is None
 
     new = schema.load(serialized)
     assert new.private_registry is False
@@ -45,6 +47,8 @@ def test_serialize_dask_environment():
     assert new.min_workers == 1
     assert new.max_workers == 2
     assert new.labels == set()
+    assert new.scheduler_spec_file is None
+    assert new.worker_spec_file is None
 
 
 def test_serialize_dask_environment_with_labels():