trainy-policy-nightly 0.1.0.dev20251111104744__py3-none-any.whl → 0.1.0.dev20251121104539__py3-none-any.whl

trainy/__init__.py

@@ -4,7 +4,7 @@ import os
 import subprocess
 
 # Replaced with the current commit when building the wheels.
-_POLICY_COMMIT_SHA = "14ff34c33b11d38c00a264197f0b885f67a65c45"
+_POLICY_COMMIT_SHA = "a7ce34b2b1ed91770b4485278b256f3dbc7d609c"
 
 
 def _get_git_commit():

trainy/policy/gke.py

@@ -10,6 +10,9 @@ kubernetes:
 
 """
 
+from copy import deepcopy
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
 import sky
 import uuid
 
@@ -22,6 +25,252 @@ DEFAULT_QUEUE = "user-queue"
 PRIORITY_CLASSES = ["low-priority", "high-priority"]
 ALLOWED_GPUS = ["H100", "H100-MEGA-80GB", "A100-80GB", "A100"]
 
+_POD_SPEC_PATH: Tuple[str, ...] = ("kubernetes", "pod_config", "spec")
+
+
+def _get_dict_path(
+    data: Dict[str, Any], path: Sequence[str]
+) -> Optional[Dict[str, Any]]:
+    """Returns the nested dict at path, creating dicts along the way."""
+    cursor: Dict[str, Any] = data
+    for key in path:
+        if key not in cursor or not isinstance(cursor[key], dict):
+            cursor[key] = {}
+        cursor = cursor[key]
+    return cursor
+
+
+def _merge_named_items(
+    base_list: Optional[List[Dict[str, Any]]],
+    override_list: Optional[List[Dict[str, Any]]],
+) -> List[Dict[str, Any]]:
+    """Merges two lists of dicts keyed by `name` while preserving order."""
+    merged: List[Dict[str, Any]] = []
+    index_by_name: Dict[str, int] = {}
+
+    def _append(item: Dict[str, Any]) -> None:
+        index_by_name[item.get("name", f"__idx_{len(merged)}")] = len(merged)
+        merged.append(item)
+
+    for entry in base_list or []:
+        if isinstance(entry, dict):
+            _append(deepcopy(entry))
+
+    for entry in override_list or []:
+        if not isinstance(entry, dict):
+            continue
+        name = entry.get("name")
+        replacement = deepcopy(entry)
+        if name is not None and name in index_by_name:
+            merged[index_by_name[name]] = replacement
+        else:
+            _append(replacement)
+
+    return merged
+
+
+def _merge_container_sections(
+    base_containers: Optional[List[Dict[str, Any]]],
+    override_containers: Optional[List[Dict[str, Any]]],
+) -> List[Dict[str, Any]]:
+    """Merges env/volumeMounts between lists of containers."""
+    base_containers = base_containers or []
+    override_containers = override_containers or []
+    merged_containers: List[Dict[str, Any]] = []
+    max_len = max(len(base_containers), len(override_containers))
+
+    for idx in range(max_len):
+        base_container = base_containers[idx] if idx < len(base_containers) else None
+        override_container = (
+            deepcopy(override_containers[idx]) if idx < len(override_containers) else None
+        )
+
+        if override_container is None:
+            if base_container is not None:
+                merged_containers.append(deepcopy(base_container))
+            continue
+
+        merged_env = _merge_named_items(
+            base_container.get("env") if base_container else None,
+            override_container.get("env"),
+        )
+        if merged_env:
+            override_container["env"] = merged_env
+        elif "env" in override_container:
+            del override_container["env"]
+
+        merged_volume_mounts = _merge_named_items(
+            base_container.get("volumeMounts") if base_container else None,
+            override_container.get("volumeMounts"),
+        )
+        if merged_volume_mounts:
+            override_container["volumeMounts"] = merged_volume_mounts
+        elif "volumeMounts" in override_container:
+            del override_container["volumeMounts"]
+
+        merged_containers.append(override_container)
+
+    return merged_containers
+
+
+def _merge_pod_spec_sections(
+    config: sky.skypilot_config.Config, override_config: Dict[str, Any]
+) -> Dict[str, Any]:
+    """Returns override_config with volumes/env/volumeMounts merged from config."""
+    merged_override = deepcopy(override_config)
+    spec = _get_dict_path(merged_override, _POD_SPEC_PATH)
+
+    base_volumes = config.get_nested(_POD_SPEC_PATH + ("volumes",), None)
+    override_volumes = spec.get("volumes")
+    merged_volumes = _merge_named_items(base_volumes, override_volumes)
+    if merged_volumes:
+        spec["volumes"] = merged_volumes
+    elif "volumes" in spec:
+        del spec["volumes"]
+
+    merged_containers = _merge_container_sections(
+        config.get_nested(_POD_SPEC_PATH + ("containers",), None),
+        spec.get("containers"),
+    )
+    if merged_containers:
+        spec["containers"] = merged_containers
+    elif "containers" in spec:
+        del spec["containers"]
+
+    merged_init_containers = _merge_container_sections(
+        config.get_nested(_POD_SPEC_PATH + ("initContainers",), None),
+        spec.get("initContainers"),
+    )
+    if merged_init_containers:
+        spec["initContainers"] = merged_init_containers
+    elif "initContainers" in spec:
+        del spec["initContainers"]
+
+    return merged_override
+
+_POD_SPEC_PATH: Tuple[str, ...] = ("kubernetes", "pod_config", "spec")
+
+
+def _get_dict_path(
+    data: Dict[str, Any], path: Sequence[str]
+) -> Optional[Dict[str, Any]]:
+    """Returns the nested dict at path, creating dicts along the way."""
+    cursor: Dict[str, Any] = data
+    for key in path:
+        if key not in cursor or not isinstance(cursor[key], dict):
+            cursor[key] = {}
+        cursor = cursor[key]
+    return cursor
+
+
+def _merge_named_items(
+    base_list: Optional[List[Dict[str, Any]]],
+    override_list: Optional[List[Dict[str, Any]]],
+) -> List[Dict[str, Any]]:
+    """Merges two lists of dicts keyed by `name` while preserving order."""
+    merged: List[Dict[str, Any]] = []
+    index_by_name: Dict[str, int] = {}
+
+    def _append(item: Dict[str, Any]) -> None:
+        index_by_name[item.get("name", f"__idx_{len(merged)}")] = len(merged)
+        merged.append(item)
+
+    for entry in base_list or []:
+        if isinstance(entry, dict):
+            _append(deepcopy(entry))
+
+    for entry in override_list or []:
+        if not isinstance(entry, dict):
+            continue
+        name = entry.get("name")
+        replacement = deepcopy(entry)
+        if name is not None and name in index_by_name:
+            merged[index_by_name[name]] = replacement
+        else:
+            _append(replacement)
+
+    return merged
+
+
+def _merge_container_sections(
+    base_containers: Optional[List[Dict[str, Any]]],
+    override_containers: Optional[List[Dict[str, Any]]],
+) -> List[Dict[str, Any]]:
+    """Merges env/volumeMounts between lists of containers."""
+    base_containers = base_containers or []
+    override_containers = override_containers or []
+    merged_containers: List[Dict[str, Any]] = []
+    max_len = max(len(base_containers), len(override_containers))
+
+    for idx in range(max_len):
+        base_container = base_containers[idx] if idx < len(base_containers) else None
+        override_container = (
+            deepcopy(override_containers[idx]) if idx < len(override_containers) else None
+        )
+
+        if override_container is None:
+            if base_container is not None:
+                merged_containers.append(deepcopy(base_container))
+            continue
+
+        merged_env = _merge_named_items(
+            base_container.get("env") if base_container else None,
+            override_container.get("env"),
+        )
+        if merged_env:
+            override_container["env"] = merged_env
+        elif "env" in override_container:
+            del override_container["env"]
+
+        merged_volume_mounts = _merge_named_items(
+            base_container.get("volumeMounts") if base_container else None,
+            override_container.get("volumeMounts"),
+        )
+        if merged_volume_mounts:
+            override_container["volumeMounts"] = merged_volume_mounts
+        elif "volumeMounts" in override_container:
+            del override_container["volumeMounts"]
+
+        merged_containers.append(override_container)
+
+    return merged_containers
+
+
+def _merge_pod_spec_sections(
+    config: sky.skypilot_config.Config, override_config: Dict[str, Any]
+) -> Dict[str, Any]:
+    """Returns override_config with volumes/env/volumeMounts merged from config."""
+    merged_override = deepcopy(override_config)
+    spec = _get_dict_path(merged_override, _POD_SPEC_PATH)
+
+    base_volumes = config.get_nested(_POD_SPEC_PATH + ("volumes",), None)
+    override_volumes = spec.get("volumes")
+    merged_volumes = _merge_named_items(base_volumes, override_volumes)
+    if merged_volumes:
+        spec["volumes"] = merged_volumes
+    elif "volumes" in spec:
+        del spec["volumes"]
+
+    merged_containers = _merge_container_sections(
+        config.get_nested(_POD_SPEC_PATH + ("containers",), None),
+        spec.get("containers"),
+    )
+    if merged_containers:
+        spec["containers"] = merged_containers
+    elif "containers" in spec:
+        del spec["containers"]
+
+    merged_init_containers = _merge_container_sections(
+        config.get_nested(_POD_SPEC_PATH + ("initContainers",), None),
+        spec.get("initContainers"),
+    )
+    if merged_init_containers:
+        spec["initContainers"] = merged_init_containers
+    elif "initContainers" in spec:
+        del spec["initContainers"]
+
+    return merged_override
+
 
 def set_tcpxo_config(user_request: sky.UserRequest) -> sky.MutatedUserRequest:
     """Sets pod specs for running TCPXO"""
@@ -32,9 +281,13 @@ def set_tcpxo_config(user_request: sky.UserRequest) -> sky.MutatedUserRequest:
             if accelerator == "H100-MEGA-80GB":
                 k8s_override_config = load_config("gke.yaml")
                 config = user_request.skypilot_config
-                new_config = sky.skypilot_config._recursive_update(
+                merged_override = _merge_pod_spec_sections(
                     config, k8s_override_config
                 )
+                new_config = sky.skypilot_config._recursive_update(
+                    config, merged_override
+                )
+
                 return sky.MutatedUserRequest(
                     task=user_request.task, skypilot_config=new_config
                 )

{trainy_policy_nightly-0.1.0.dev20251111104744.dist-info → trainy_policy_nightly-0.1.0.dev20251121104539.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: trainy-policy-nightly
-Version: 0.1.0.dev20251111104744
+Version: 0.1.0.dev20251121104539
 Summary: Trainy Skypilot Policy
 Author: Andrew Aikawa
 Author-email: asai@berkeley.edu

{trainy_policy_nightly-0.1.0.dev20251111104744.dist-info → trainy_policy_nightly-0.1.0.dev20251121104539.dist-info}/RECORD

@@ -1,4 +1,4 @@
-trainy/__init__.py,sha256=JOzvZYiDSOsMOMFV3WUJWII72wYWqtturfqtMNJJTUA,1236
+trainy/__init__.py,sha256=-YRjM9Yuq-e3fxZkRsqp_7BgvmNxsdBgmg7MZlhYJzk,1236
 trainy/config.py,sha256=MkAVzIvbS61y73uneWzZYS_FikGYZHI__DuPmRJ1lH0,322
 trainy/configs/eks.yaml,sha256=QlUFr3w7ZcTU0c0jqoGxUlWN1nVl5yE_b5IWd0ToDgs,554
 trainy/configs/gke.yaml,sha256=aB52IPKZhvb2ZoNn32RMkP299z3PoW2CVE-mHA5j-oE,4972
@@ -6,8 +6,8 @@ trainy/configs/nebius.yaml,sha256=hH0W_EKIYe5EaRC9pkUHWX3P7hbM46tg0OLrc6QatJw,44
 trainy/logging.py,sha256=CMHdnnZ2b-rNDTfU7bi6jQ3AmxVv0M4sOlghnVJIqko,1186
 trainy/policy/__init__.py,sha256=DElGvZKzOqFbUs90Pi53ZpNBE3BFsalgpYMos_4tboU,180
 trainy/policy/eks.py,sha256=FBgRprnSjmp94xQvCV0tnKTgxKBk2SPNIVNFg6PPrZI,2493
-trainy/policy/gke.py,sha256=ClkVT6W_cWosiW1dhnJCfun_euE0YSy8B5UA-u-7f9I,8173
+trainy/policy/gke.py,sha256=1OLEt9g9vrt3bZOHuA3v86xzgBFsxh5yHP3InKW4mh0,16967
 trainy/policy/nebius.py,sha256=79zYL0Z6JsOomB0IF3CimSnkZfyyPvNYaJNQb0iv_Uw,3062
-trainy_policy_nightly-0.1.0.dev20251111104744.dist-info/METADATA,sha256=50fw9N5uKazKy_M0bs4Bv40sm_NZy8gQW1W9edjrN7E,1228
-trainy_policy_nightly-0.1.0.dev20251111104744.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-trainy_policy_nightly-0.1.0.dev20251111104744.dist-info/RECORD,,
+trainy_policy_nightly-0.1.0.dev20251121104539.dist-info/METADATA,sha256=InPLPXltF-kjnkPFaS28QkKHpT1uwhUiw8Vhzm5DveM,1228
+trainy_policy_nightly-0.1.0.dev20251121104539.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+trainy_policy_nightly-0.1.0.dev20251121104539.dist-info/RECORD,,