Read worker identity and rpc port from pod metadata

The serving processes take workers per pod, pods per cell, gpu slots and the rpc port from MILES_WORKER_METADATA instead of the spec they rebuilt.

miles/utils/workers/serving/serve.py

@@ -18,7 +18,7 @@ def main() -> None:
     _log(f"start own_argv={own_argv} worker_argv={worker_argv}")
 
     spec = compute_serve_worker_spec(specs_fn=args.specs, pool_id=args.pool_id, worker_argv=worker_argv)
-    identity = read_worker_identity(scheduling=spec.scheduling, environ=os.environ)
+    identity = read_worker_identity(os.environ)
     env_vars = spec.env_var(
         WorkerLaunchContext(
             args=spec.args,

miles/utils/workers/serving/serve_inner.py

@@ -19,7 +19,11 @@ from miles.utils.workers.serving.utils import (
     parse_own_args,
     split_worker_argv,
 )
-from miles.utils.workers.serving.worker_identity import read_worker_identity, read_worker_in_pod_index
+from miles.utils.workers.serving.worker_identity import (
+    read_worker_identity,
+    read_worker_in_pod_index,
+    read_worker_metadata,
+)
 from miles.utils.workers.types import ClusterBackend
 from miles.utils.workers.worker_spec import RPC_PORT_NAME, BaseServeSpec, PortInfo
 
@@ -41,7 +45,7 @@ def main() -> None:
 
 
 def create_worker(spec: BaseServeSpec, *, specs_fn: str, worker_argv: list[str]) -> Any:
-    identity = read_worker_identity(scheduling=spec.scheduling, environ=os.environ)
+    identity = read_worker_identity(os.environ)
     _log(f"identity={identity}")
     capability = DeferredBackendCapability(create=lambda: _backend_capability(specs_fn, worker_argv))
     context = identity.ctor_context(args=spec.args, capability=capability)
@@ -57,7 +61,7 @@ def _backend_capability(specs_fn: str, worker_argv: list[str]) -> BackendCapabil
 
 
 def _rpc_port_of(spec: BaseServeSpec) -> PortInfo:
-    ports = [port_info for port_info in spec.port_infos if port_info.name == RPC_PORT_NAME]
+    ports = [port_info for port_info in read_worker_metadata(os.environ).port_infos if port_info.name == RPC_PORT_NAME]
     assert len(ports) == 1, f"spec '{spec.name}' declares {len(ports)} rpc ports, so this process cannot pick one"
     return ports[0]
 

miles/utils/workers/serving/worker_identity.py

@@ -5,8 +5,14 @@ from dataclasses import dataclass
 from typing import Any
 
 from miles.utils.workers.backend_capability.base import BackendCapability
-from miles.utils.workers.env_vars import CELL_INDEX_ENV_VAR, POD_INDEX_ENV_VAR, SUBPROCESS_INDEX_ENV_VAR
-from miles.utils.workers.worker_spec import SchedulingSpec, WorkerCtorContext
+from miles.utils.workers.connection_config import WorkerPodMetadata
+from miles.utils.workers.env_vars import (
+    CELL_INDEX_ENV_VAR,
+    POD_INDEX_ENV_VAR,
+    SUBPROCESS_INDEX_ENV_VAR,
+    WORKER_METADATA_ENV_VAR,
+)
+from miles.utils.workers.worker_spec import WorkerCtorContext
 
 
 @dataclass(frozen=True)
@@ -36,9 +42,11 @@ class KubernetesWorkerIdentity:
         )
 
 
-def read_worker_identity(*, scheduling: SchedulingSpec, environ: Mapping[str, str]) -> KubernetesWorkerIdentity:
-    workers_per_pod = scheduling.workers_per_pod()
-    pods_per_cell = scheduling.pods_per_cell()
+def read_worker_identity(environ: Mapping[str, str]) -> KubernetesWorkerIdentity:
+    metadata = read_worker_metadata(environ)
+    workers_per_pod = metadata.workers_per_pod
+    pods_per_cell = metadata.pods_per_cell
+    gpu_slots_per_worker = metadata.gpu_slots_per_worker
 
     worker_in_pod_index = read_worker_in_pod_index(
         environ,
@@ -73,19 +81,19 @@ def read_worker_identity(*, scheduling: SchedulingSpec, environ: Mapping[str, st
         CELL_INDEX_ENV_VAR,
         required_because="nothing else tells this pod which cell of its pool it belongs to",
     )
-    assert (
-        cell_index < scheduling.num_cells
-    ), f"{CELL_INDEX_ENV_VAR} is {cell_index}, but the pool is scheduled with {scheduling.num_cells} cells"
-
     return KubernetesWorkerIdentity(
         cell_index=cell_index,
         pod_in_cell_index=pod_in_cell_index,
         worker_in_pod_index=worker_in_pod_index,
         workers_per_pod=workers_per_pod,
-        gpu_slots_per_worker=scheduling.num_gpu_slots_per_worker,
+        gpu_slots_per_worker=gpu_slots_per_worker,
     )
 
 
+def read_worker_metadata(environ: Mapping[str, str]) -> WorkerPodMetadata:
+    return WorkerPodMetadata.model_validate_json(environ[WORKER_METADATA_ENV_VAR])
+
+
 def read_worker_in_pod_index(environ: Mapping[str, str], *, required_because: str | None = None) -> int:
     return _index_from(environ, SUBPROCESS_INDEX_ENV_VAR, required_because=required_because)