Add env vars support to Pyfunc ensembler

api/api/openapi.bundle.yaml

@@ -1951,6 +1951,11 @@ components:
               max_replica: 6
               memory_request: memory_request
               cpu_request: cpu_request
+            env:
+            - name: name
+              value: value
+            - name: name
+              value: value
             timeout: timeout
           updated_at: 2000-01-23T04:56:07.000+00:00
           standard_config:
@@ -2212,6 +2217,11 @@ components:
             max_replica: 6
             memory_request: memory_request
             cpu_request: cpu_request
+          env:
+          - name: name
+            value: value
+          - name: name
+            value: value
           timeout: timeout
         updated_at: 2000-01-23T04:56:07.000+00:00
         standard_config:
@@ -2347,6 +2357,11 @@ components:
           max_replica: 6
           memory_request: memory_request
           cpu_request: cpu_request
+        env:
+        - name: name
+          value: value
+        - name: name
+          value: value
         timeout: timeout
       nullable: true
       properties:
@@ -2359,6 +2374,10 @@ components:
         timeout:
           pattern: ^[0-9]+(ms|s|m|h)$
           type: string
+        env:
+          items:
+            $ref: '#/components/schemas/EnvVar'
+          type: array
       required:
       - ensembler_id
       - project_id
@@ -2502,6 +2521,11 @@ components:
                 max_replica: 6
                 memory_request: memory_request
                 cpu_request: cpu_request
+              env:
+              - name: name
+                value: value
+              - name: name
+                value: value
               timeout: timeout
             updated_at: 2000-01-23T04:56:07.000+00:00
             standard_config:
@@ -2634,6 +2658,11 @@ components:
               max_replica: 6
               memory_request: memory_request
               cpu_request: cpu_request
+            env:
+            - name: name
+              value: value
+            - name: name
+              value: value
             timeout: timeout
           updated_at: 2000-01-23T04:56:07.000+00:00
           standard_config:

api/api/specs/routers.yaml

@@ -13,7 +13,7 @@ info:
 .timeout: &timeout
   type: "string"
   pattern: '^[0-9]+(ms|s|m|h)$'
-  
+
 paths:
   "/projects/{project_id}/routers":
     get:
@@ -851,6 +851,10 @@ components:
           $ref: "#/components/schemas/ResourceRequest"
         timeout:
           <<: *timeout
+        env:
+          type: "array"
+          items:
+            $ref: "common.yaml#/components/schemas/EnvVar"
 
     ResourceRequest:
       type: "object"

api/db-migrations/000010_add_pyfunc_config.down.sql

@@ -1,4 +1,4 @@
--- Removes py_func_ref_config from table.
+-- Removes pyfunc_config from table.
 -- Hence, this migration involves data loss for ensemblers with type that is "pyfunc"
 ALTER TABLE ensembler_configs DROP COLUMN pyfunc_config;
 

api/db-migrations/000010_add_pyfunc_config.up.sql

@@ -1,3 +1,3 @@
--- Adds py_func_ref_config to table.
+-- Adds pyfunc_config to table.
 ALTER TABLE ensembler_configs ADD pyfunc_config jsonb;
 

api/db-migrations/000011_add_pyfunc_config_env.down.sql

@@ -0,0 +1,9 @@
+-- Migrate pyfunc_config data from new schema to old schema (remove the env column)
+UPDATE ensembler_configs
+SET pyfunc_config = (SELECT json_build_object(
+        'timeout', pyfunc_config ->> 'timeout',
+        'project_id', (pyfunc_config ->> 'project_id')::int,
+        'ensembler_id', (pyfunc_config ->> 'ensembler_id')::int,
+        'resource_request', (pyfunc_config ->> 'resource_request')::jsonb
+    ))
+WHERE type ='pyfunc';

api/db-migrations/000011_add_pyfunc_config_env.up.sql

@@ -0,0 +1,10 @@
+-- Migrate pyfunc_config data from old schema to new schema (with new 'env' field set as an empty array by default)
+UPDATE ensembler_configs
+SET pyfunc_config = (SELECT json_build_object(
+        'env', '[]'::jsonb,
+        'timeout', pyfunc_config ->> 'timeout',
+        'project_id', (pyfunc_config ->> 'project_id')::int,
+        'ensembler_id', (pyfunc_config ->> 'ensembler_id')::int,
+        'resource_request', (pyfunc_config ->> 'resource_request')::jsonb
+    ))
+WHERE type ='pyfunc';

api/turing/api/request/request.go

@@ -148,7 +148,7 @@ func (r RouterConfig) BuildRouterVersion(
 		}
 		if r.Ensembler.Type == models.EnsemblerPyFuncType {
 			if r.Ensembler.PyfuncConfig == nil {
-				return nil, errors.New("missing ensembler pyfunc reference config")
+				return nil, errors.New("missing ensembler pyfunc config")
 			}
 
 			// Verify if the ensembler given by its ProjectID and EnsemblerID exist

api/turing/models/ensembler.go

@@ -58,6 +58,8 @@ type EnsemblerPyfuncConfig struct {
 	ResourceRequest *ResourceRequest `json:"resource_request" validate:"required"`
 	// Request timeout in duration format e.g. 60s
 	Timeout string `json:"timeout" validate:"required"`
+	// Environment variables to set in the container
+	Env EnvVars `json:"env" validate:"required"`
 }
 
 type ExperimentMapping struct {

api/turing/service/router_deployment_service.go

@@ -455,7 +455,7 @@ func (ds *deploymentService) buildEnsemblerServiceImage(
 		Timeout:         routerVersion.Ensembler.PyfuncConfig.Timeout,
 		Endpoint:        PyFuncEnsemblerServiceEndpoint,
 		Port:            PyFuncEnsemblerServicePort,
-		Env:             models.EnvVars{},
+		Env:             routerVersion.Ensembler.PyfuncConfig.Env,
 		ServiceAccount:  "",
 	}
 

sdk/samples/router/edit_a_router.py

@@ -47,17 +47,14 @@ def main(turing_api: str, project: str):
     new_router_config_to_save = router_1.config
     new_router_config_to_save.resource_request.min_replica = 0
 
-    router_1.create_version(new_router_config_to_save)
+    new_undeployed_version = turing.router.config.router_version.RouterVersion.create(new_router_config_to_deploy,
+                                                                                      router_id=1)
 
     # Notice that the latest router version is undeployed (Turing has created the new version without deploying it)
     versions = router_1.list_versions()
     for v in versions:
         print(v)
 
-    # Alternatively, you may also do the following to create a new version for your router without deploying it.
-    new_undeployed_version = turing.router.config.router_version.RouterVersion.create(new_router_config_to_save,
-                                                                                      router_id=1)
-
 
 if __name__ == '__main__':
     import fire

sdk/tests/conftest.py

@@ -334,12 +334,13 @@ def generic_ensembler_docker_config(generic_resource_request, generic_env_var):
 
 
 @pytest.fixture
-def generic_ensembler_pyfunc_config(generic_resource_request):
+def generic_ensembler_pyfunc_config(generic_resource_request, generic_env_var):
     return turing.generated.models.EnsemblerPyfuncConfig(
         project_id=77,
         ensembler_id=11,
         resource_request=generic_resource_request,
-        timeout="500ms"
+        timeout="500ms",
+        env=[generic_env_var]
     )
 
 

sdk/tests/router/config/router_ensembler_config_test.py

@@ -66,8 +66,9 @@ def test_create_router_ensembler_config(id, type, standard_config, docker_config
     ).to_open_api()
     assert actual == request.getfixturevalue(expected)
 
+
 @pytest.mark.parametrize(
-    "project_id,ensembler_id,resource_request,timeout,expected", [
+    "project_id,ensembler_id,resource_request,timeout,env,expected", [
         pytest.param(
             77,
             11,
@@ -78,6 +79,11 @@ def test_create_router_ensembler_config(id, type, standard_config, docker_config
                 memory_request='512Mi'
             ),
             "500ms",
+            [
+                EnvVar(
+                    name="env_name",
+                    value="env_val")
+            ],
             "generic_pyfunc_router_ensembler_config"
         )
     ])
@@ -86,6 +92,7 @@ def test_create_pyfunc_router_ensembler_config(
         ensembler_id,
         resource_request,
         timeout,
+        env,
         expected,
         request
 ):
@@ -94,12 +101,13 @@ def test_create_pyfunc_router_ensembler_config(
         ensembler_id=ensembler_id,
         resource_request=resource_request,
         timeout=timeout,
+        env=env
     ).to_open_api()
     assert actual == request.getfixturevalue(expected)
 
 
 @pytest.mark.parametrize(
-    "project_id,ensembler_id,resource_request,timeout,expected", [
+    "project_id,ensembler_id,resource_request,timeout,env,expected", [
         pytest.param(
             77,
             11,
@@ -110,6 +118,11 @@ def test_create_pyfunc_router_ensembler_config(
                 memory_request='512Mi'
             ),
             "500ks",
+            [
+                EnvVar(
+                    name="env_name",
+                    value="env_val")
+            ],
             ApiValueError
         )
     ])
@@ -118,6 +131,7 @@ def test_create_pyfunc_router_ensembler_config_with_invalid_timeout(
         ensembler_id,
         resource_request,
         timeout,
+        env,
         expected
 ):
     with pytest.raises(expected):
@@ -126,6 +140,7 @@ def test_create_pyfunc_router_ensembler_config_with_invalid_timeout(
             ensembler_id=ensembler_id,
             resource_request=resource_request,
             timeout=timeout,
+            env=env
         ).to_open_api()
 
 

sdk/turing/generated/model/ensembler_pyfunc_config.py

@@ -27,7 +27,9 @@
 )
 
 def lazy_import():
+    from turing.generated.model.env_var import EnvVar
     from turing.generated.model.resource_request import ResourceRequest
+    globals()['EnvVar'] = EnvVar
     globals()['ResourceRequest'] = ResourceRequest
 
 
@@ -86,6 +88,7 @@ def openapi_types():
             'ensembler_id': (int,),  # noqa: E501
             'resource_request': (ResourceRequest,),  # noqa: E501
             'timeout': (str,),  # noqa: E501
+            'env': ([EnvVar],),  # noqa: E501
         }
 
     @cached_property
@@ -98,6 +101,7 @@ def discriminator():
         'ensembler_id': 'ensembler_id',  # noqa: E501
         'resource_request': 'resource_request',  # noqa: E501
         'timeout': 'timeout',  # noqa: E501
+        'env': 'env',  # noqa: E501
     }
 
     _composed_schemas = {}
@@ -152,6 +156,7 @@ def __init__(self, project_id, ensembler_id, resource_request, timeout, *args, *
                                 Animal class but this time we won't travel
                                 through its discriminator because we passed in
                                 _visited_composed_classes = (Animal,)
+            env ([EnvVar]): [optional]  # noqa: E501
         """
 
         _check_type = kwargs.pop('_check_type', True)

sdk/turing/router/config/router_ensembler_config.py

@@ -99,6 +99,7 @@ def pyfunc_config(self, pyfunc_config: turing.generated.models.EnsemblerPyfuncCo
         elif isinstance(pyfunc_config, dict):
             pyfunc_config["resource_request"] = \
                 turing.generated.models.ResourceRequest(**pyfunc_config["resource_request"])
+            pyfunc_config["env"] = [turing.generated.models.EnvVar(**env_var) for env_var in pyfunc_config["env"]]
             self._pyfunc_config = turing.generated.models.EnsemblerPyfuncConfig(
                 **pyfunc_config
             )
@@ -127,19 +128,22 @@ def __init__(self,
                  project_id: int,
                  ensembler_id: int,
                  timeout: str,
-                 resource_request: ResourceRequest):
+                 resource_request: ResourceRequest,
+                 env: List['EnvVar']):
         """
         Method to create a new Pyfunc ensembler
 
         :param project_id: project id of the current project
         :param ensembler_id: ensembler_id of the ensembler
         :param resource_request: ResourceRequest instance containing configs related to the resources required
         :param timeout: request timeout which when exceeded, the request to the ensembler will be terminated
+        :param env: environment variables required by the container
         """
         self.project_id = project_id
         self.ensembler_id = ensembler_id
         self.resource_request = resource_request
         self.timeout = timeout
+        self.env = env
         super().__init__(type="pyfunc")
 
     @property
@@ -174,12 +178,23 @@ def timeout(self) -> str:
     def timeout(self, timeout: str):
         self._timeout = timeout
 
+    @property
+    def env(self) -> List['EnvVar']:
+        return self._env
+
+    @env.setter
+    def env(self, env: List['EnvVar']):
+        self._env = env
+
     def to_open_api(self) -> OpenApiModel:
+        assert all(isinstance(env_var, EnvVar) for env_var in self.env)
+
         self.pyfunc_config = turing.generated.models.EnsemblerPyfuncConfig(
             project_id=self.project_id,
             ensembler_id=self.ensembler_id,
             resource_request=self.resource_request.to_open_api(),
             timeout=self.timeout,
+            env=[env_var.to_open_api() for env_var in self.env],
         )
         return super().to_open_api()
 

sdk/turing/router/config/router_version.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import turing
-import turing.generated.models
 import dataclasses
 from enum import Enum
 

ui/src/router/components/configuration/components/pyfunc_config_section/PyFuncRefConfigTable.js → ui/src/router/components/configuration/components/pyfunc_config_section/PyFuncConfigTable.js

@@ -2,7 +2,7 @@ import React, { useContext } from "react";
 import { EuiDescriptionList } from "@elastic/eui";
 import EnsemblersContext from "../../../../../providers/ensemblers/context";
 
-export const PyFuncRefConfigTable = ({ config: { ensembler_id } }) => {
+export const PyFuncConfigTable = ({ config: { ensembler_id } }) => {
   const { ensemblers } = useContext(EnsemblersContext);
 
   const ensembler_name = Object.values(ensemblers)