openmpi: support custom resources.

kubeflow · May 9, 2018 · 664e5a3 · 664e5a3
1 parent 18b670d
commit 664e5a3
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/kubeflow/openmpi/prototypes/openmpi.jsonnet b/kubeflow/openmpi/prototypes/openmpi.jsonnet
@@ -13,6 +13,7 @@
 // @optionalParam gpu number 0 Number of GPUs per worker.
 // @optionalParam cpu string null CPU limits per worker.
 // @optionalParam memory string null Memory limits per worker.
+// @optionalParam customResources string null Comma-delimited list of "resourceName=amount" pairs which you want to limit per worker.
 // @optionalParam serviceAccountName string null the service account name to run pods. The service account should have clusterRoleBinding for "view" ClusterRole.  If it was not set, service account and its role binding will be created automatically.
 // @optionalParam schedulerName string default-scheduler scheduler name to use for the components.
 // @optionalParam controllerImage string jiez/openmpi-controller:0.0.1 Docker image of the openmpi-controller.

diff --git a/kubeflow/openmpi/workloads.libsonnet b/kubeflow/openmpi/workloads.libsonnet
@@ -99,7 +99,7 @@ local ROLE_WORKER = "worker";
       name: "openmpi-job",
       image: params.image,
       imagePullPolicy: params.imagePullPolicy,
-      resources: $.resources(params, role),
+      resources: std.mergePatch($.resources(params, role), $.customResources(params, role)),
       terminationMessagePath: "/dev/termination-log",
       terminationMessagePolicy: "File",
       workingDir: "/kubeflow/openmpi/data",
@@ -173,4 +173,9 @@ local ROLE_WORKER = "worker";
 
   nodeSelector(params, role)::
     if role == ROLE_WORKER then util.toObject(params.nodeSelector) else {},
+
+  customResources(params, role)::
+    if role == ROLE_WORKER then {
+      limits: util.toObject(params.customResources),
+    } else {},
 }