Add missing env vars to TaskAPI

pkg/workloads/k8s.go

@@ -657,12 +657,12 @@ func tensorFlowHandlerContainers(api *spec.API, envVars []kcore.EnvVar, isJob bo
 func taskEnvVars(api *spec.API) []kcore.EnvVar {
 	envVars := apiContainerEnvVars(api)
 	envVars = append(envVars,
-
 		kcore.EnvVar{
 			Name:  "CORTEX_TASK_SPEC",
 			Value: TaskSpecPath,
 		},
 	)
+	envVars = append(envVars, getKubexitEnvVars(APIContainerName)...)
 	return envVars
 }
 

test/e2e/e2e/tests.py

@@ -183,12 +183,16 @@ def test_batch_api(
         job_spec = response.json()
 
         # monitor job progress
+        job_id = job_spec["job_id"]
+        endpoint_override = (
+            f"http://localhost:8888/batch/{api_name}?jobID={job_id}" if local_operator else None
+        )
         assert job_done(
             client=client,
-            api_name=job_spec["api_name"],
-            job_id=job_spec["job_id"],
+            api_name=api_name,
+            job_id=job_id,
             timeout=job_timeout,
-            local_operator=local_operator,
+            endpoint_override=endpoint_override,
         ), f"job did not succeed (api_name: {api_name}, job_id: {job_spec['job_id']})"
 
     except:
@@ -330,6 +334,7 @@ def test_task_api(
     job_timeout: int = None,
     retry_attempts: int = 0,
     api_config_name: str = "cortex.yaml",
+    local_operator: bool = False,
 ):
     api_dir = TEST_APIS_DIR / api
     with open(str(api_dir / api_config_name)) as f:
@@ -341,28 +346,34 @@ def test_task_api(
     client.create_api(api_spec=api_specs[0], project_dir=str(api_dir))
 
     try:
+        endpoint_override = f"http://localhost:8888/tasks/{api_name}" if local_operator else None
         assert endpoint_ready(
-            client=client, api_name=api_name, timeout=deploy_timeout
+            client=client,
+            api_name=api_name,
+            timeout=deploy_timeout,
+            endpoint_override=endpoint_override,
         ), f"api {api_name} not ready"
 
         response = None
         for _ in range(retry_attempts + 1):
-            response = request_task(
-                client,
-                api_name,
-            )
+            response = request_task(client, api_name, local_operator=local_operator)
             if response.status_code == HTTPStatus.OK:
                 break
 
             time.sleep(1)
 
         job_spec = response.json()
 
+        job_id = job_spec["job_id"]
+        endpoint_override = (
+            f"http://localhost:8888/tasks/{api_name}?jobID={job_id}" if local_operator else None
+        )
         assert job_done(
             client=client,
             api_name=api_name,
             job_id=job_spec["job_id"],
             timeout=job_timeout,
+            endpoint_override=endpoint_override,
         ), f"task job did not succeed (api_name: {api_name}, job_id: {job_spec['job_id']})"
 
     except:

test/e2e/e2e/utils.py

@@ -96,11 +96,15 @@ def _is_ready():
 
 
 def job_done(
-    client: cx.Client, api_name: str, job_id: str, timeout: int = None, local_operator: bool = False
+    client: cx.Client,
+    api_name: str,
+    job_id: str,
+    timeout: int = None,
+    endpoint_override: str = None,
 ) -> bool:
     def _is_ready():
-        if local_operator:
-            job_info = requests.get(f"http://localhost:8888/batch/{api_name}?jobID={job_id}")
+        if endpoint_override:
+            job_info = requests.get(endpoint_override)
             job_info = job_info.json()
             return job_info["job_status"]["status"] == "succeeded"
 
@@ -211,9 +215,13 @@ def request_task(
     api_name: str,
     config: Dict = None,
     timeout: int = None,
+    local_operator: bool = False,
 ):
-    api_info = client.get_api(api_name)
-    endpoint = api_info["endpoint"]
+    if local_operator:
+        endpoint = f"http://localhost:8888/tasks/{api_name}"
+    else:
+        api_info = client.get_api(api_name)
+        endpoint = api_info["endpoint"]
 
     payload = {}
     if config is not None:

test/e2e/tests/aws/test_task.py

@@ -32,4 +32,5 @@ def test_task_api(printer: Callable, config: Dict, client: cx.Client, api: str):
         retry_attempts=5,
         deploy_timeout=config["global"]["task_deploy_timeout"],
         job_timeout=config["global"]["task_job_timeout"],
+        local_operator=config["global"]["local_operator"],
     )