Repository: PrefectHQ/prefect

Issue:
Add executor kwargs to `KubernetesJobEnvironment`
I recently realized that to configure a `KubernetesJobEnvironment` to use a non-default executor requires knowledge of the appropriate environment variables.  We should expose a better interface for this (similar to the `RemoteEnvironment`)


Discussion:
We should extend this to also add the option to the `FargateTaskEnvironment`

Resolving patch:
diff --git a/src/prefect/environments/execution/fargate/fargate_task.py b/src/prefect/environments/execution/fargate/fargate_task.py
--- a/src/prefect/environments/execution/fargate/fargate_task.py
+++ b/src/prefect/environments/execution/fargate/fargate_task.py
@@ -60,6 +60,8 @@ class FargateTaskEnvironment(Environment):
             `AWS_SESSION_TOKEN` or `None`
         - region_name (str, optional): AWS region name for connecting the boto3 client.
             Defaults to the value set in the environment variable `REGION_NAME` or `None`
+        - executor_kwargs (dict, optional): a dictionary of kwargs to be passed to
+            the executor; defaults to an empty dictionary
         - labels (List[str], optional): a list of labels, which are arbitrary string identifiers used by Prefect
             Agents when polling for work
         - on_start (Callable, optional): a function callback which will be called before the flow begins to run
@@ -74,6 +76,7 @@ def __init__(  # type: ignore
         aws_secret_access_key: str = None,
         aws_session_token: str = None,
         region_name: str = None,
+        executor_kwargs: dict = None,
         labels: List[str] = None,
         on_start: Callable = None,
         on_exit: Callable = None,
@@ -90,6 +93,7 @@ def __init__(  # type: ignore
         # Parse accepted kwargs for definition and run
         self.task_definition_kwargs, self.task_run_kwargs = self._parse_kwargs(kwargs)
 
+        self.executor_kwargs = executor_kwargs or dict()
         super().__init__(labels=labels, on_start=on_start, on_exit=on_exit)
 
     def _parse_kwargs(self, user_kwargs: dict) -> tuple:
@@ -311,7 +315,7 @@ def run_flow(self) -> None:
                 flow = cloudpickle.load(f)
 
                 runner_cls = get_default_flow_runner_class()
-                executor_cls = get_default_executor_class()()
+                executor_cls = get_default_executor_class()(**self.executor_kwargs)
                 runner_cls(flow=flow).run(executor=executor_cls)
         except Exception as exc:
             self.logger.exception(
diff --git a/src/prefect/environments/execution/k8s/job.py b/src/prefect/environments/execution/k8s/job.py
--- a/src/prefect/environments/execution/k8s/job.py
+++ b/src/prefect/environments/execution/k8s/job.py
@@ -39,6 +39,8 @@ class KubernetesJobEnvironment(Environment):
 
     Args:
         - job_spec_file (str, optional): Path to a job spec YAML file
+        - executor_kwargs (dict, optional): a dictionary of kwargs to be passed to
+            the executor; defaults to an empty dictionary
         - labels (List[str], optional): a list of labels, which are arbitrary string identifiers used by Prefect
             Agents when polling for work
         - on_start (Callable, optional): a function callback which will be called before the flow begins to run
@@ -48,11 +50,13 @@ class KubernetesJobEnvironment(Environment):
     def __init__(
         self,
         job_spec_file: str = None,
+        executor_kwargs: dict = None,
         labels: List[str] = None,
         on_start: Callable = None,
         on_exit: Callable = None,
     ) -> None:
         self.job_spec_file = os.path.abspath(job_spec_file) if job_spec_file else None
+        self.executor_kwargs = executor_kwargs or dict()
 
         # Load specs from file if path given, store on object
         self._job_spec = self._load_spec_from_file()
@@ -161,7 +165,7 @@ def run_flow(self) -> None:
                 flow = cloudpickle.load(f)
 
                 runner_cls = get_default_flow_runner_class()
-                executor_cls = get_default_executor_class()()
+                executor_cls = get_default_executor_class()(**self.executor_kwargs)
                 runner_cls(flow=flow).run(executor=executor_cls)
         except Exception as exc:
             self.logger.exception(


Test patch:
diff --git a/tests/environments/execution/test_fargate_task_environment.py b/tests/environments/execution/test_fargate_task_environment.py
--- a/tests/environments/execution/test_fargate_task_environment.py
+++ b/tests/environments/execution/test_fargate_task_environment.py
@@ -15,12 +15,19 @@
 def test_create_fargate_task_environment():
     environment = FargateTaskEnvironment()
     assert environment
+    assert environment.executor_kwargs == {}
     assert environment.labels == set()
     assert environment.on_start is None
     assert environment.on_exit is None
     assert environment.logger.name == "prefect.FargateTaskEnvironment"
 
 
+def test_create_fargate_task_environment_with_executor_kwargs():
+    environment = FargateTaskEnvironment(executor_kwargs={"test": "here"})
+    assert environment
+    assert environment.executor_kwargs == {"test": "here"}
+
+
 def test_create_fargate_task_environment_labels():
     environment = FargateTaskEnvironment(labels=["foo"])
     assert environment
@@ -401,7 +408,7 @@ def test_execute_run_task_agent_token(monkeypatch):
 
 
 def test_run_flow(monkeypatch):
-    environment = FargateTaskEnvironment()
+    environment = FargateTaskEnvironment(executor_kwargs={"test": "here"})
 
     flow_runner = MagicMock()
     monkeypatch.setattr(
@@ -409,6 +416,11 @@ def test_run_flow(monkeypatch):
         MagicMock(return_value=flow_runner),
     )
 
+    executor = MagicMock()
+    monkeypatch.setattr(
+        "prefect.engine.get_default_executor_class", MagicMock(return_value=executor),
+    )
+
     with tempfile.TemporaryDirectory() as directory:
         with open(os.path.join(directory, "flow_env.prefect"), "w+"):
             flow = prefect.Flow("test")
@@ -423,6 +435,7 @@ def test_run_flow(monkeypatch):
                 environment.run_flow()
 
         assert flow_runner.call_args[1]["flow"].name == "test"
+        assert executor.call_args[1] == {"test": "here"}
 
 
 def test_run_flow_calls_callbacks(monkeypatch):
diff --git a/tests/environments/execution/test_k8s_job_environment.py b/tests/environments/execution/test_k8s_job_environment.py
--- a/tests/environments/execution/test_k8s_job_environment.py
+++ b/tests/environments/execution/test_k8s_job_environment.py
@@ -25,12 +25,27 @@ def test_create_k8s_job_environment():
         )
         assert environment
         assert environment.job_spec_file == os.path.join(directory, "job.yaml")
+        assert environment.executor_kwargs == {}
         assert environment.labels == set()
         assert environment.on_start is None
         assert environment.on_exit is None
         assert environment.logger.name == "prefect.KubernetesJobEnvironment"
 
 
+def test_create_k8s_job_environment_with_executor_kwargs():
+    with tempfile.TemporaryDirectory() as directory:
+
+        with open(os.path.join(directory, "job.yaml"), "w+") as file:
+            file.write("job")
+
+        environment = KubernetesJobEnvironment(
+            job_spec_file=os.path.join(directory, "job.yaml"),
+            executor_kwargs={"test": "here"},
+        )
+        assert environment
+        assert environment.executor_kwargs == {"test": "here"}
+
+
 def test_create_k8s_job_environment_labels():
     with tempfile.TemporaryDirectory() as directory:
 
@@ -195,7 +210,9 @@ def test_create_flow_run_job_fails_outside_cluster():
 
 def test_run_flow(monkeypatch):
     file_path = os.path.dirname(prefect.environments.execution.dask.k8s.__file__)
-    environment = KubernetesJobEnvironment(path.join(file_path, "job.yaml"))
+    environment = KubernetesJobEnvironment(
+        path.join(file_path, "job.yaml"), executor_kwargs={"test": "here"}
+    )
 
     flow_runner = MagicMock()
     monkeypatch.setattr(
@@ -203,6 +220,11 @@ def test_run_flow(monkeypatch):
         MagicMock(return_value=flow_runner),
     )
 
+    executor = MagicMock()
+    monkeypatch.setattr(
+        "prefect.engine.get_default_executor_class", MagicMock(return_value=executor),
+    )
+
     with tempfile.TemporaryDirectory() as directory:
         with open(os.path.join(directory, "flow_env.prefect"), "w+") as env:
             flow = prefect.Flow("test")
@@ -217,6 +239,7 @@ def test_run_flow(monkeypatch):
                 environment.run_flow()
 
         assert flow_runner.call_args[1]["flow"].name == "test"
+        assert executor.call_args[1] == {"test": "here"}
 
 
 def test_run_flow_calls_callbacks(monkeypatch):