Repository: skypilot-org/skypilot
Pull request: [Serve] Support env var in service yaml
<!-- Describe the changes in this PR -->
Fixes https://github.com/skypilot-org/skypilot/issues/3048

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [x] Any manual or new tests for this PR (please specify below)
   - [x] Run sky serve with environment variable used in `readiness_probe` in service yaml (e.g. `readiness_probe: /${VERSION}/models`)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
[Serve] Support env var in `service` field
When creating an example for serving llm, it will be much easier if the service field can support the following:

```yaml
service:
  readiness_probe:
    path: /v1/chat/completions
    post_data:
      model: $MODEL_NAME
      messages:
        - role: user
          content: How to print hello world?

envs:
  MODEL_NAME: codellama/CodeLlama-70b-Python-hf

run: |
  python -u -m vllm.entrypoints.openai.api_server \
                --host 0.0.0.0 \
                --model $MODEL_NAME \
                --tensor-parallel-size $SKYPILOT_NUM_GPUS_PER_NODE | tee ~/openai_api_server.log
```
cc'ing @cblmemo @MaoZiming

Hints from issue discussion:
Agree! would be great to have

Commit message: [serve] add env vars to smoke test

Patch to review:
diff --git a/sky/task.py b/sky/task.py
index 36de66370d0..17da10c95b1 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -72,11 +72,11 @@ def _is_valid_name(name: Optional[str]) -> bool:
     return bool(re.fullmatch(_VALID_NAME_REGEX, name))
 
 
-def _fill_in_env_vars_in_file_mounts(
-    file_mounts: Dict[str, Any],
+def _fill_in_env_vars(
+    yaml_field: Dict[str, Any],
     task_envs: Dict[str, str],
 ) -> Dict[str, Any]:
-    """Detects env vars in file_mounts and fills them with task_envs.
+    """Detects env vars in yaml field and fills them with task_envs.
 
     Use cases of env vars in file_mounts:
     - dst/src paths; e.g.,
@@ -84,8 +84,19 @@ def _fill_in_env_vars_in_file_mounts(
     - storage's name (bucket name)
     - storage's source (local path)
 
-    We simply dump file_mounts into a json string, and replace env vars using
-    regex. This should be safe as file_mounts has been schema-validated.
+    Use cases of env vars in service:
+    - model type; e.g.,
+        service:
+          readiness_probe:
+            path: /v1/chat/completions
+            post_data:
+              model: $MODEL_NAME
+              messages:
+                - role: user
+                  content: How to print hello world?
+
+    We simply dump yaml_field into a json string, and replace env vars using
+    regex. This should be safe as yaml config has been schema-validated.
 
     Env vars of the following forms are detected:
         - ${ENV}
@@ -93,7 +104,7 @@ def _fill_in_env_vars_in_file_mounts(
     where <ENV> must appear in task.envs.
     """
     # TODO(zongheng): support ${ENV:-default}?
-    file_mounts_str = json.dumps(file_mounts)
+    yaml_field_str = json.dumps(yaml_field)
 
     def replace_var(match):
         var_name = match.group(1)
@@ -102,8 +113,8 @@ def replace_var(match):
 
     # Pattern for valid env var names in bash.
     pattern = r'\$\{?\b([a-zA-Z_][a-zA-Z0-9_]*)\b\}?'
-    file_mounts_str = re.sub(pattern, replace_var, file_mounts_str)
-    return json.loads(file_mounts_str)
+    yaml_field_str = re.sub(pattern, replace_var, yaml_field_str)
+    return json.loads(yaml_field_str)
 
 
 def _check_docker_login_config(task_envs: Dict[str, str]) -> bool:
@@ -366,8 +377,13 @@ def from_yaml_config(
         # Fill in any Task.envs into file_mounts (src/dst paths, storage
         # name/source).
         if config.get('file_mounts') is not None:
-            config['file_mounts'] = _fill_in_env_vars_in_file_mounts(
-                config['file_mounts'], config.get('envs', {}))
+            config['file_mounts'] = _fill_in_env_vars(config['file_mounts'],
+                                                      config.get('envs', {}))
+
+        # Fill in any Task.envs into service (e.g. MODEL_NAME).
+        if config.get('service') is not None:
+            config['service'] = _fill_in_env_vars(config['service'],
+                                                  config.get('envs', {}))
 
         task = Task(
             config.pop('name', None),
diff --git a/tests/skyserve/llm/service.yaml b/tests/skyserve/llm/service.yaml
index 48b5344db59..02f494d69b6 100644
--- a/tests/skyserve/llm/service.yaml
+++ b/tests/skyserve/llm/service.yaml
@@ -1,7 +1,10 @@
 service:
-  readiness_probe: /v1/models
+  readiness_probe: /${API_VERSION}/models
   replicas: 1
 
+envs:
+  API_VERSION: v1
+
 resources:
   ports: 8087
   cloud: gcp


Code review comments:
diff_hunk: @@ -1,7 +1,10 @@
 service:
-  readiness_probe: /v1/models
+  readiness_probe: /${API_VERSION}/models
path: tests/skyserve/llm/service.yaml
text: @user1:
It will be great to use sth similar to the original issue (#3048) in our smoke test, as it is more widely used and could test env vars in service/run section in the same time. Considering use `MODEL_NAME` for `lmsys/fastchat-t5-3b-v1.0`

@author:
Good suggestion, just updated!