fix: actually support add/drop capabilities in structs (#2236)

determined-ai · Apr 15, 2021 · 0b2c2c7 · 0b2c2c7
1 parent 1c926e0
commit 0b2c2c7
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 0 deletions.
diff --git a/harness/determined/common/schemas/expconf/_v0.py b/harness/determined/common/schemas/expconf/_v0.py
@@ -295,6 +295,8 @@ def __init__(
 
 class EnvironmentConfigV0(schemas.SchemaBase):
     _id = "http://determined.ai/schemas/expconf/v0/environment.json"
+    add_capabilities: Optional[List[str]] = None
+    drop_capabilities: Optional[List[str]] = None
     environment_variables: Optional[EnvironmentVariablesV0] = None
     force_pull_image: Optional[bool] = None
     image: Optional[EnvironmentImageV0] = None
@@ -305,6 +307,8 @@ class EnvironmentConfigV0(schemas.SchemaBase):
     @schemas.auto_init
     def __init__(
         self,
+        add_capabilities: Optional[List[str]] = None,
+        drop_capabilities: Optional[List[str]] = None,
         environment_variables: Optional[EnvironmentVariablesV0] = None,
         force_pull_image: Optional[bool] = None,
         image: Optional[EnvironmentImageV0] = None,

diff --git a/master/pkg/schemas/expconf/environment_config.go b/master/pkg/schemas/expconf/environment_config.go
@@ -22,6 +22,9 @@ type EnvironmentConfigV0 struct {
 	RegistryAuth   *types.AuthConfig `json:"registry_auth"`
 	ForcePullImage *bool             `json:"force_pull_image"`
 	PodSpec        *k8sV1.Pod        `json:"pod_spec"`
+
+	AddCapabilities  *[]string `json:"add_capabilities"`
+	DropCapabilities *[]string `json:"drop_capabilities"`
 }
 
 //go:generate ../gen.sh

diff --git a/master/pkg/schemas/expconf/zgen_environment_config_v0.go b/master/pkg/schemas/expconf/zgen_environment_config_v0.go
diff --git a/schemas/test_cases/v0/experiment.yaml b/schemas/test_cases/v0/experiment.yaml
@@ -45,6 +45,10 @@
         serveraddress: srvaddr
         identitytoken: idtkn
         registrytoken: rgtkn
+      add_capabilities:
+        - CAP_CHOWN
+      drop_capabilities:
+        - CAP_KILL
     hyperparameters:
       global_batch_size:
         type: const
@@ -157,6 +161,8 @@
       pod_spec:
       ports: {}
       registry_auth: null
+      add_capabilities: []
+      drop_capabilities: []
     hyperparameters:
       global_batch_size:
         type: const