add env & env_from spec

kubeflow · Oct 31, 2023 · 3c5ad4a · 3c5ad4a
1 parent d2e311f
commit 3c5ad4a
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/cmd/suggestion/optuna/v1beta1/requirements.txt b/cmd/suggestion/optuna/v1beta1/requirements.txt
@@ -2,3 +2,4 @@ grpcio>=1.41.1
 protobuf>=3.19.5, <=3.20.3
 googleapis-common-protos==1.53.0
 optuna>=3.0.0
+cmaes>=0.10.0
diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py
@@ -25,6 +25,7 @@
 from kubeflow.katib.constants import constants
 from kubeflow.katib.utils import utils
 from kubernetes import client, config
+from kubernetes.client.models import V1EnvVar, V1EnvFromSource
 
 
 class KatibClient(object):
@@ -140,6 +141,8 @@ def tune(
         parameters: Dict[str, Any],
         base_image: str = constants.BASE_IMAGE_TENSORFLOW,
         namespace: Optional[str] = None,
+        env: Union[Dict[str, str], List[V1EnvVar], None] = None,
+        env_from: Union[V1EnvFromSource, List[V1EnvFromSource], None] = None,
         algorithm_name: str = "random",
         algorithm_settings: Union[dict, List[models.V1beta1AlgorithmSetting], None] = None,
         objective_metric_name: str = None,
@@ -172,6 +175,13 @@ def tune(
                 objective function.
             base_image: Image to use when executing the objective function.
             namespace: Namespace for the Experiment.
+            env: Environment variable(s) to be attached to each trial container. You can either specifiy
+                a list of kubernetes.client.models.V1EnvVar (documented here:
+                https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1EnvVar.md) or a dictionary  
+                corresponding to the environment variable name and value pair(s).
+            env_from: Source(s) of environment variables to be populated in each trial container. You can either specify
+                a kubernetes.client.models.V1EnvFromSource (documented here:
+                https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1EnvFromSource.md) or a list of such a type.
             algorithm_name: Search algorithm for the HyperParameter tuning.
             algorithm_settings: Settings for the search algorithm given.
                 For available fields, check this doc: https://www.kubeflow.org/docs/components/katib/experiment/#search-algorithms-in-detail.
@@ -319,6 +329,12 @@ def tune(
                 limits=resources_per_trial,
             )
 
+        if isinstance(env, dict):
+            env = [V1EnvVar(name=str(k), value=str(v)) for k, v in env.items()]
+
+        if isinstance(env_from, V1EnvFromSource):
+            env_from = [env_from]
+
         # Create Trial specification.
         trial_spec = client.V1Job(
             api_version="batch/v1",
@@ -336,6 +352,8 @@ def tune(
                                 image=base_image,
                                 command=["bash", "-c"],
                                 args=[exec_script],
+                                env=env,
+                                env_from=env_from,
                                 resources=resources_per_trial,
                             )
                         ],