Merge pull request PrefectHQ#3046 from emcake/k8s-customize-yaml

joshmeek · web-flow · commit 766deb6b30b3 · 2020-08-25T14:37:49.000-04:00
allow customization of the yaml template
diff --git a/changes/pr3046.yaml b/changes/pr3046.yaml
@@ -0,0 +1,5 @@
+enhancement:
+  - "Allow for setting path to a custom job YAML spec on the Kubernetes Agent - [#3046](https://github.com/PrefectHQ/prefect/pull/3046)"
+
+contributor:
+  - "[emcake](https://github.com/emcake)"
diff --git a/src/prefect/agent/kubernetes/agent.py b/src/prefect/agent/kubernetes/agent.py
@@ -157,6 +157,8 @@ def replace_job_spec_yaml(self, flow_run: GraphQLResult, image: str) -> dict:
                 container registry, such as Amazon ECR.
         - `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
@@ -165,7 +167,10 @@ def replace_job_spec_yaml(self, flow_run: GraphQLResult, image: str) -> dict:
         Returns:
             - dict: a dictionary representation of a k8s job for flow execution
         """
-        with open(path.join(path.dirname(__file__), "job_spec.yaml"), "r") as job_file:
+        yaml_path = os.getenv(
+            "YAML_TEMPLATE", path.join(path.dirname(__file__), "job_spec.yaml")
+        )
+        with open(yaml_path, "r") as job_file:
             job = yaml.safe_load(job_file)
 
         identifier = str(uuid.uuid4())[:8]
