Add resources to K8sContainerContext (#7619)

Summary: The next step after this would be to add these as fields on the instance/helm chart and in the  executor so that they can be configured on the instance and per-op, but this gives us what we need as a foundation to vary them per-location and have that config threaded all the way through the launched k8s jobs for that location.

docs/content/deployment/guides/kubernetes/deploying-with-helm.mdx

@@ -292,6 +292,13 @@ dagster-user-deployments:
         - name: test-volume
           mountPath: /opt/dagster/test_folder
           subPath: test_file.yaml
+      resources:
+        limits:
+          cpu: 100m
+          memory: 128Mi
+        requests:
+          cpu: 100m
+          memory: 128Mi
       includeConfigInLaunchedRuns:
         enabled: true
 ```

examples/docs_snippets/docs_snippets/deploying/kubernetes/user_code_deployment_config.yaml

@@ -23,5 +23,12 @@ dagster-user-deployments:
         - name: test-volume
           mountPath: /opt/dagster/test_folder
           subPath: test_file.yaml
+      resources:
+        limits:
+          cpu: 100m
+          memory: 128Mi
+        requests:
+          cpu: 100m
+          memory: 128Mi
       includeConfigInLaunchedRuns:
         enabled: true

helm/dagster/charts/dagster-user-deployments/templates/helpers/_helpers.tpl

@@ -122,6 +122,9 @@ DAGSTER_K8S_PIPELINE_RUN_ENV_CONFIGMAP: "{{ template "dagster.fullname" . }}-pip
     {{- if .labels }}
     labels: {{- .labels | toYaml | nindent 6 }}
     {{- end }}
+    {{- if .resources }}
+    resources: {{- .resources | toYaml | nindent 6 }}
+    {{- end }}
     namespace: {{ $.Release.Namespace }}
     service_account_name: {{ include "dagsterUserDeployments.serviceAccountName" $ }}
   {{- end }}

helm/dagster/schema/schema/charts/dagster/subschema/run_launcher.py

@@ -64,6 +64,7 @@ class K8sRunLauncherConfig(BaseModel):
     volumes: List[kubernetes.Volume]
     labels: Optional[Dict[str, str]]
     failPodOnRunFailure: Optional[bool]
+    resources: Optional[kubernetes.ResourceRequirements]
 
     class Config:
         extra = Extra.forbid

helm/dagster/schema/schema/charts/utils/kubernetes.py

@@ -149,3 +149,8 @@ class Config:
 class Volume(BaseModel):
     class Config:
         schema_extra = {"$ref": create_definition_ref("io.k8s.api.core.v1.Volume")}
+
+
+class ResourceRequirements(BaseModel):
+    class Config:
+        schema_extra = {"$ref": create_definition_ref("io.k8s.api.core.v1.ResourceRequirements")}

helm/dagster/schema/schema_tests/test_instance.py

@@ -191,6 +191,36 @@ def test_k8s_run_launcher_fail_pod_on_run_failure(template: HelmTemplate):
     assert run_launcher_config["config"]["fail_pod_on_run_failure"]
 
 
+def test_k8s_run_launcher_resources(template: HelmTemplate):
+    resources = {
+        "requests": {"memory": "64Mi", "cpu": "250m"},
+        "limits": {"memory": "128Mi", "cpu": "500m"},
+    }
+
+    helm_values = DagsterHelmValues.construct(
+        runLauncher=RunLauncher.construct(
+            type=RunLauncherType.K8S,
+            config=RunLauncherConfig.construct(
+                k8sRunLauncher=K8sRunLauncherConfig.construct(
+                    imagePullPolicy="Always",
+                    loadInclusterConfig=True,
+                    envConfigMaps=[],
+                    envSecrets=[],
+                    envVars=[],
+                    volumeMounts=[],
+                    volumes=[],
+                    resources=resources,
+                )
+            ),
+        )
+    )
+    configmaps = template.render(helm_values)
+    instance = yaml.full_load(configmaps[0].data["dagster.yaml"])
+    run_launcher_config = instance["run_launcher"]
+
+    assert run_launcher_config["config"]["resources"] == resources
+
+
 def test_celery_k8s_run_launcher_config(template: HelmTemplate):
     image = {"repository": "test_repo", "tag": "test_tag", "pullPolicy": "Always"}
 

helm/dagster/schema/schema_tests/test_user_deployments.py

@@ -758,6 +758,54 @@ def test_user_deployment_labels(template: HelmTemplate, include_config_in_launch
         _assert_no_container_context(user_deployments[0])
 
 
+@pytest.mark.parametrize("include_config_in_launched_runs", [False, True])
+def test_user_deployment_resources(template: HelmTemplate, include_config_in_launched_runs: bool):
+    name = "foo"
+
+    resources = {
+        "requests": {"memory": "64Mi", "cpu": "250m"},
+        "limits": {"memory": "128Mi", "cpu": "500m"},
+    }
+
+    deployment = UserDeployment.construct(
+        name=name,
+        image={"repository": f"repo/{name}", "tag": "tag1", "pullPolicy": "Always"},
+        dagsterApiGrpcArgs=["-m", name],
+        port=3030,
+        resources=resources,
+        includeConfigInLaunchedRuns={"enabled": include_config_in_launched_runs},
+    )
+
+    helm_values = DagsterHelmValues.construct(
+        dagsterUserDeployments=UserDeployments(
+            enabled=True,
+            enableSubchart=True,
+            deployments=[deployment],
+        )
+    )
+
+    user_deployments = template.render(helm_values)
+
+    assert len(user_deployments) == 1
+
+    if include_config_in_launched_runs:
+        container_context = user_deployments[0].spec.template.spec.containers[0].env[2]
+        assert container_context.name == "DAGSTER_CLI_API_GRPC_CONTAINER_CONTEXT"
+        assert json.loads(container_context.value) == {
+            "k8s": {
+                "image_pull_policy": "Always",
+                "env_config_maps": [
+                    "release-name-dagster-user-deployments-foo-user-env",
+                ],
+                "resources": resources,
+                "namespace": "default",
+                "service_account_name": "release-name-dagster-user-deployments-user-deployments",
+            }
+        }
+    else:
+        _assert_no_container_context(user_deployments[0])
+
+
 def test_subchart_image_pull_secrets(subchart_template: HelmTemplate):
     image_pull_secrets = [{"name": "super-duper-secret"}]
     deployment_values = DagsterUserDeploymentsHelmValues.construct(

helm/dagster/templates/helpers/instance/_run-launcher.tpl

@@ -120,6 +120,10 @@ config:
   labels: {{- $k8sRunLauncherConfig.labels | toYaml | nindent 4 }}
   {{- end }}
 
+  {{- if $k8sRunLauncherConfig.resources }}
+  resources: {{- $k8sRunLauncherConfig.resources | toYaml | nindent 4 }}
+  {{- end }}
+
   {{- if $k8sRunLauncherConfig.failPodOnRunFailure }}
   fail_pod_on_run_failure: true
   {{- end }}

helm/dagster/values.yaml

@@ -470,6 +470,19 @@ runLauncher:
       #   my_label_key: my_label_value
       labels: {}
 
+      # Default compute resource requirements for the container in the Job's Pod. See:
+      # https://kubernetes.io/docs/concepts/configuration/manage-resources-containers
+      #
+      # Example:
+      # resources:
+      #   limits:
+      #     cpu: 100m
+      #     memory: 128Mi
+      #   requests:
+      #     cpu: 100m
+      #     memory: 128Mi
+      resources: {}
+
       # Whether the launched Kubernetes Jobs and Pods should fail if the Dagster run fails.
       failPodOnRunFailure: false
 

python_modules/libraries/dagster-k8s/dagster_k8s/container_context.py

@@ -33,6 +33,7 @@ class K8sContainerContext(
             ("volumes", List[Dict[str, Any]]),
             ("labels", Dict[str, str]),
             ("namespace", Optional[str]),
+            ("resources", Dict[str, Any]),
         ],
     )
 ):
@@ -53,6 +54,7 @@ def __new__(
         volumes: Optional[List[Dict[str, Any]]] = None,
         labels: Optional[Dict[str, str]] = None,
         namespace: Optional[str] = None,
+        resources: Optional[Dict[str, Any]] = None,
     ):
         return super(K8sContainerContext, cls).__new__(
             cls,
@@ -72,6 +74,7 @@ def __new__(
             ],
             labels=check.opt_dict_param(labels, "labels"),
             namespace=check.opt_str_param(namespace, "namespace"),
+            resources=check.opt_dict_param(resources, "resources"),
         )
 
     def merge(self, other: "K8sContainerContext") -> "K8sContainerContext":
@@ -94,6 +97,7 @@ def merge(self, other: "K8sContainerContext") -> "K8sContainerContext":
             volumes=_dedupe_list(other.volumes + self.volumes),
             labels=merge_dicts(other.labels, self.labels),
             namespace=other.namespace if other.namespace else self.namespace,
+            resources=other.resources if other.resources else self.resources,
         )
 
     @staticmethod
@@ -115,6 +119,7 @@ def create_for_run(
                     volumes=run_launcher.volumes,
                     labels=run_launcher.labels,
                     namespace=run_launcher.job_namespace,
+                    resources=run_launcher.resources,
                 )
             )
 
@@ -162,6 +167,7 @@ def create_from_config(run_container_context) -> "K8sContainerContext":
             volumes=processed_context_value.get("volumes"),
             labels=processed_context_value.get("labels"),
             namespace=processed_context_value.get("namespace"),
+            resources=processed_context_value.get("resources"),
         )
 
     def get_k8s_job_config(self, job_image, run_launcher) -> DagsterK8sJobConfig:
@@ -179,4 +185,5 @@ def get_k8s_job_config(self, job_image, run_launcher) -> DagsterK8sJobConfig:
             volume_mounts=self.volume_mounts,
             volumes=self.volumes,
             labels=self.labels,
+            resources=self.resources,
         )

python_modules/libraries/dagster-k8s/dagster_k8s/executor.py

@@ -84,6 +84,7 @@ def k8s_job_executor(init_context: InitExecutorContext) -> Executor:
         volumes=exc_cfg.get("volumes"),
         labels=exc_cfg.get("labels"),
         namespace=exc_cfg.get("job_namespace"),
+        resources=exc_cfg.get("resources"),
     )
 
     return StepDelegatingExecutor(

python_modules/libraries/dagster-k8s/dagster_k8s/job.py

@@ -211,7 +211,7 @@ class DagsterK8sJobConfig(
         "_K8sJobTaskConfig",
         "job_image dagster_home image_pull_policy image_pull_secrets service_account_name "
         "instance_config_map postgres_password_secret env_config_maps env_secrets env_vars "
-        "volume_mounts volumes labels",
+        "volume_mounts volumes labels resources",
     )
 ):
     """Configuration parameters for launching Dagster Jobs on Kubernetes.
@@ -255,6 +255,8 @@ class DagsterK8sJobConfig(
             https://v1-18.docs.kubernetes.io/docs/reference/generated/kubernetes-api/v1.18/#volume-v1-core
         labels (Optional[Dict[str, str]]): Additional labels that should be included in the Job's Pod. See:
             https://kubernetes.io/docs/concepts/overview/working-with-objects/labels
+        resources (Optional[Dict[str, Any]]) Compute resource requirements for the container. See:
+            https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/
     """
 
     def __new__(
@@ -272,6 +274,7 @@ def __new__(
         volume_mounts=None,
         volumes=None,
         labels=None,
+        resources=None,
     ):
         return super(DagsterK8sJobConfig, cls).__new__(
             cls,
@@ -300,6 +303,7 @@ def __new__(
                 for volume in check.opt_list_param(volumes, "volumes")
             ],
             labels=check.opt_dict_param(labels, "labels", key_type=str, value_type=str),
+            resources=check.opt_dict_param(resources, "resources", key_type=str),
         )
 
     @classmethod
@@ -456,6 +460,17 @@ def config_type_container(cls):
                 description="Labels to apply to all created pods. See: "
                 "https://kubernetes.io/docs/concepts/overview/working-with-objects/labels",
             ),
+            "resources": Field(
+                Noneable(
+                    {
+                        "limits": Field(dict, is_required=False),
+                        "requests": Field(dict, is_required=False),
+                    }
+                ),
+                is_required=False,
+                description="Compute resource requirements for the container. See: "
+                "https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/",
+            ),
         }
 
     @classmethod
@@ -607,6 +622,8 @@ def construct_dagster_k8s_job(
 
     user_defined_k8s_volume_mounts = container_config.pop("volume_mounts", [])
 
+    user_defined_resources = container_config.pop("resources", {})
+
     volume_mounts = (
         [
             {
@@ -621,6 +638,8 @@ def construct_dagster_k8s_job(
         + user_defined_k8s_volume_mounts
     )
 
+    resources = user_defined_resources if user_defined_resources else job_config.resources
+
     container_config = merge_dicts(
         container_config,
         {
@@ -631,6 +650,7 @@ def construct_dagster_k8s_job(
             "env": env + job_config.env + additional_k8s_env_vars,
             "env_from": job_config.env_from_sources + additional_k8s_env_from,
             "volume_mounts": volume_mounts,
+            "resources": resources,
         },
     )
 

python_modules/libraries/dagster-k8s/dagster_k8s/launcher.py

@@ -66,6 +66,7 @@ def __init__(
         volumes=None,
         labels=None,
         fail_pod_on_run_failure=None,
+        resources=None,
     ):
         self._inst_data = check.opt_inst_param(inst_data, "inst_data", ConfigurableClassData)
         self.job_namespace = check.str_param(job_namespace, "job_namespace")
@@ -109,6 +110,7 @@ def __init__(
         self._fail_pod_on_run_failure = check.opt_bool_param(
             fail_pod_on_run_failure, "fail_pod_on_run_failure"
         )
+        self._resources = check.opt_dict_param(resources, "resources")
 
         super().__init__()
 
@@ -144,6 +146,10 @@ def volume_mounts(self):
     def volumes(self):
         return self._volumes
 
+    @property
+    def resources(self):
+        return self._resources
+
     @property
     def env_vars(self):
         return self._env_vars

python_modules/libraries/dagster-k8s/dagster_k8s_tests/unit_tests/conftest.py

@@ -29,6 +29,12 @@ def kubeconfig_file(tmp_path):
     return str(config_path)
 
 
+LAUNCHER_RESOURCES = {
+    "requests": {"cpu": "128m", "memory": "64Mi"},
+    "limits": {"cpu": "500m", "memory": "1000Mi"},
+}
+
+
 @pytest.fixture
 def k8s_run_launcher_instance(kubeconfig_file):  # pylint: disable=redefined-outer-name
     with environ({"BAR_TEST": "bar"}):
@@ -46,6 +52,7 @@ def k8s_run_launcher_instance(kubeconfig_file):  # pylint: disable=redefined-out
                         "load_incluster_config": False,
                         "kubeconfig_file": kubeconfig_file,
                         "env_vars": ["BAR_TEST"],
+                        "resources": LAUNCHER_RESOURCES,
                     },
                 },
             }