Source code for torchx.schedulers.kubernetes_scheduler
#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
# pyre-strict
"""
This contains the TorchX Kubernetes scheduler which can be used to run TorchX
components on a Kubernetes cluster.
Prerequisites
==============
The TorchX Kubernetes scheduler depends on Volcano. If you're trying to do an
upgrade you'll need to completely remove all non-Job Volcano resources and recreate.
Install Volcano:
.. code:: bash
kubectl apply -f https://raw.githubusercontent.com/volcano-sh/volcano/v1.6.0/installer/volcano-development.yaml
See the
`Volcano Quickstart <https://github.com/volcano-sh/volcano>`_
for more information.
Pod Overlay
===========
You can overlay arbitrary Kubernetes Pod fields on generated pods by setting
the ``kubernetes`` metadata on your role. The value can be:
- A dict with the overlay structure
- A resource URI pointing to a YAML file (e.g. ``file://``, ``s3://``, ``gs://``)
Merge semantics:
- **dict**: recursive merge (upsert)
- **list**: append by default, replace if tuple (Python) or ``!!python/tuple`` tag (YAML)
- **primitives**: replace
.. code:: python
from torchx.specs import Role
# Dict overlay - lists append, tuples replace
role = Role(
name="trainer",
image="my-image:latest",
entrypoint="train.py",
metadata={
"kubernetes": {
"spec": {
"nodeSelector": {"gpu": "true"},
"tolerations": [{"key": "nvidia.com/gpu", "operator": "Exists"}], # appends
"volumes": ({"name": "my-volume", "emptyDir": {}},) # replaces
}
}
}
)
# File URI overlay
role = Role(
name="trainer",
image="my-image:latest",
entrypoint="train.py",
metadata={
"kubernetes": "file:///path/to/pod_overlay.yaml"
}
)
CLI usage with builtin components:
.. code:: bash
$ torchx run --scheduler kubernetes dist.ddp \\
--metadata kubernetes=file:///path/to/pod_overlay.yaml \\
--script train.py
Example ``pod_overlay.yaml``:
.. code:: yaml
spec:
nodeSelector:
node.kubernetes.io/instance-type: p4d.24xlarge
tolerations:
- key: nvidia.com/gpu
operator: Exists
effect: NoSchedule
volumes: !!python/tuple
- name: my-volume
emptyDir: {}
The overlay is deep-merged with the generated pod, preserving existing fields
and adding or overriding specified ones.
"""
import json
import logging
import re
import warnings
from dataclasses import dataclass
from datetime import datetime
from typing import Any, cast, Iterable, Mapping, TYPE_CHECKING
import torchx
import yaml
from torchx.schedulers.api import (
DescribeAppResponse,
filter_regex,
ListAppResponse,
Scheduler,
split_lines,
Stream,
StructuredOpts,
)
from torchx.schedulers.ids import make_unique
from torchx.specs.api import (
AppDef,
AppDryRunInfo,
AppState,
BindMount,
CfgVal,
DeviceMount,
macros,
ReplicaState,
ReplicaStatus,
RetryPolicy,
Role,
RoleStatus,
runopts,
VolumeMount,
)
from torchx.specs.overlays import apply_overlay, get_overlay
from torchx.util.colors import BLUE, ENDC
from torchx.util.strings import normalize_str
from torchx.workspace.docker_workspace import DockerWorkspaceMixin
if TYPE_CHECKING:
from docker import DockerClient
from kubernetes.client import ApiClient, CustomObjectsApi
from kubernetes.client.models import ( # noqa: F401 imported but unused
V1Container,
V1Pod,
)
from kubernetes.client.rest import ApiException
logger: logging.Logger = logging.getLogger(__name__)
# Kubernetes reserves a small amount of resources per host for the system. For
# TorchX we always assume the entire host is being requested so we adjust the
# requested numbers account for the node reserved resources.
#
# https://kubernetes.io/docs/tasks/administer-cluster/reserve-compute-resources/
RESERVED_MILLICPU = 100
RESERVED_MEMMB = 1024
def _apply_pod_overlay(
pod: "V1Pod",
overlay: dict[str, Any],
) -> None:
"""Apply overlay dict to V1Pod object, merging nested fields.
Uses :py:func:`~torchx.specs.overlays.apply_overlay` with operator support:
- Default: dicts merge recursively, lists append, primitives overwrite.
- Use :py:func:`~torchx.specs.overlays.PUT` to replace a value entirely.
- Use :py:func:`~torchx.specs.overlays.JOIN` for strategic merge of list
items by key field.
- Use :py:func:`~torchx.specs.overlays.DEL` to remove a key.
.. note:: Only ``pod.spec`` and ``pod.metadata`` are updated from the
merged result. Other top-level V1Pod fields (e.g., ``apiVersion``,
``kind``, ``status``) in the overlay are applied during merging but
not copied back to the pod object.
"""
from kubernetes import client
api = client.ApiClient()
pod_dict = api.sanitize_for_serialization(pod)
apply_overlay(pod_dict, overlay)
# pyrefly: ignore [missing-attribute]
merged_pod = api._ApiClient__deserialize(pod_dict, "V1Pod")
pod.spec = merged_pod.spec
pod.metadata = merged_pod.metadata
RETRY_POLICIES: Mapping[str, Iterable[Mapping[str, str]]] = {
RetryPolicy.REPLICA: [],
RetryPolicy.APPLICATION: [
{"event": "PodEvicted", "action": "RestartJob"},
{"event": "PodFailed", "action": "RestartJob"},
],
}
JOB_STATE: dict[str, AppState] = {
# Pending is the phase that job is pending in the queue, waiting for
# scheduling decision
"Pending": AppState.PENDING,
# Aborting is the phase that job is aborted, waiting for releasing pods
"Aborting": AppState.RUNNING,
# Aborted is the phase that job is aborted by user or error handling
"Aborted": AppState.CANCELLED,
# Running is the phase that minimal available tasks of Job are running
"Running": AppState.RUNNING,
# Restarting is the phase that the Job is restarted, waiting for pod
# releasing and recreating
"Restarting": AppState.RUNNING,
# Completed is the phase that all tasks of Job are completed successfully
"Completed": AppState.SUCCEEDED,
# Completing is the phase that the Job is in the process of completing
"Completing": AppState.RUNNING,
# Terminating is the phase that the Job is terminated, waiting for releasing
# pods
"Terminating": AppState.RUNNING,
# Teriminated is the phase that the job is finished unexpected, e.g. events
"Terminated": AppState.FAILED,
"Failed": ReplicaState.FAILED,
}
TASK_STATE: dict[str, ReplicaState] = {
# Pending means the task is pending in the apiserver.
"Pending": ReplicaState.PENDING,
# Allocated means the scheduler assigns a host to it.
"Allocated": ReplicaState.PENDING,
# Pipelined means the scheduler assigns a host to wait for releasing
# resource.
"Pipelined": ReplicaState.PENDING,
# Binding means the scheduler send Bind request to apiserver.
"Binding": ReplicaState.PENDING,
# Bound means the task/Pod bounds to a host.
"Bound": ReplicaState.PENDING,
# Running means a task is running on the host.
"Running": ReplicaState.RUNNING,
# Releasing means a task/pod is deleted.
"Releasing": ReplicaState.RUNNING,
# Succeeded means that all containers in the pod have voluntarily
# terminated with a container exit code of 0, and the system is not
# going to restart any of these containers.
"Succeeded": ReplicaState.SUCCEEDED,
# Failed means that all containers in the pod have terminated, and at
# least one container has terminated in a failure (exited with a
# non-zero exit code or was stopped by the system).
"Failed": ReplicaState.FAILED,
# Unknown means the status of task/pod is unknown to the scheduler.
"Unknown": ReplicaState.UNKNOWN,
}
LABEL_VERSION = "torchx.pytorch.org/version"
LABEL_APP_NAME = "torchx.pytorch.org/app-name"
LABEL_ROLE_INDEX = "torchx.pytorch.org/role-index"
LABEL_ROLE_NAME = "torchx.pytorch.org/role-name"
LABEL_REPLICA_ID = "torchx.pytorch.org/replica-id"
LABEL_KUBE_APP_NAME = "app.kubernetes.io/name"
LABEL_ORGANIZATION = "app.kubernetes.io/managed-by"
LABEL_UNIQUE_NAME = "app.kubernetes.io/instance"
ANNOTATION_ISTIO_SIDECAR = "sidecar.istio.io/inject"
LABEL_INSTANCE_TYPE = "node.kubernetes.io/instance-type"
def prefix_container_name(container_name: str, role_name: str, replica_id: int) -> str:
"""
Generate a prefix for a container name.
Returns empty string for default container (role_name-replica_id), name for others.
"""
default_container = f"{role_name}-{replica_id}"
if container_name == default_container:
return ""
return f"{BLUE}{container_name}{ENDC} "
# role.env translates to static env variables in the yaml
# {"FOO" : "bar"} =====> - name: FOO
# value: bar
# unless this placeholder is present at the start of the role.env value then the env variable
# in the yaml will be dynamically populated at runtime (placeholder is stripped out of the value)
# {"FOO" : "[FIELD_PATH]bar"} =====> - name: FOO
# valueFrom:
# fieldRef:
# fieldPath: bar
PLACEHOLDER_FIELD_PATH = "[FIELD_PATH]"
[docs]
def sanitize_for_serialization(obj: object) -> object:
from kubernetes import client
api = client.ApiClient()
return api.sanitize_for_serialization(obj)
[docs]
def role_to_pod(
name: str,
role: Role,
service_account: str | None,
reserved_millicpu: int = RESERVED_MILLICPU,
reserved_memmb: int = RESERVED_MEMMB,
efa_device_count: int | None = None,
) -> "V1Pod":
from kubernetes.client.models import ( # noqa: F811 redefinition of unused
V1Container,
V1ContainerPort,
V1EmptyDirVolumeSource,
V1EnvVar,
V1EnvVarSource,
V1HostPathVolumeSource,
V1ObjectFieldSelector,
V1ObjectMeta,
V1PersistentVolumeClaimVolumeSource,
V1Pod,
V1PodSpec,
V1ResourceRequirements,
V1SecurityContext,
V1Volume,
V1VolumeMount,
)
# limits puts an upper cap on the resources a pod may consume.
# requests is how much the scheduler allocates. We assume that the jobs will
# be allocation the whole machine so requests is slightly lower than the
# requested resources to account for the Kubernetes node reserved resources.
limits = {}
requests = {}
resource = role.resource
if resource.cpu > 0:
mcpu = int(resource.cpu * 1000)
limits["cpu"] = f"{mcpu}m"
request_mcpu = max(mcpu - reserved_millicpu, 0)
requests["cpu"] = f"{request_mcpu}m"
if resource.memMB > 0:
limits["memory"] = f"{int(resource.memMB)}M"
request_memMB = max(int(resource.memMB) - reserved_memmb, 0)
requests["memory"] = f"{request_memMB}M"
if resource.gpu > 0:
requests["nvidia.com/gpu"] = limits["nvidia.com/gpu"] = str(resource.gpu)
EFA_DEVICE = "vpc.amazonaws.com/efa"
for device_name, device_limit in resource.devices.items():
limits[device_name] = str(device_limit)
# Handle EFA device count override:
# - None (default): use whatever count is in the resource spec (already added above)
# - 0: remove EFA devices entirely
# - N > 0: set EFA device count to N (override or add)
if efa_device_count is not None:
if efa_device_count == 0:
limits.pop(EFA_DEVICE, None)
else:
limits[EFA_DEVICE] = str(efa_device_count)
resources = V1ResourceRequirements(
limits=limits,
requests=requests,
)
node_selector: dict[str, str] = {}
if LABEL_INSTANCE_TYPE in resource.capabilities:
node_selector[LABEL_INSTANCE_TYPE] = resource.capabilities[LABEL_INSTANCE_TYPE]
# To support PyTorch dataloaders we need to set /dev/shm to larger than the
# 64M default so we mount an unlimited sized tmpfs directory on it.
SHM_VOL = "dshm"
volumes = [
V1Volume(
name=SHM_VOL,
empty_dir=V1EmptyDirVolumeSource(
medium="Memory",
),
),
]
volume_mounts = [
V1VolumeMount(name=SHM_VOL, mount_path="/dev/shm"),
]
security_context = V1SecurityContext()
for i, mount in enumerate(role.mounts):
mount_name = f"mount-{i}"
if isinstance(mount, BindMount):
volumes.append(
V1Volume(
name=mount_name,
host_path=V1HostPathVolumeSource(
path=mount.src_path,
),
)
)
volume_mounts.append(
V1VolumeMount(
name=mount_name,
mount_path=mount.dst_path,
read_only=mount.read_only,
)
)
elif isinstance(mount, VolumeMount):
volumes.append(
V1Volume(
name=mount_name,
persistent_volume_claim=V1PersistentVolumeClaimVolumeSource(
claim_name=mount.src,
),
)
)
volume_mounts.append(
V1VolumeMount(
name=mount_name,
mount_path=mount.dst_path,
read_only=mount.read_only,
)
)
elif isinstance(mount, DeviceMount):
volumes.append(
V1Volume(
name=mount_name,
host_path=V1HostPathVolumeSource(
path=mount.src_path,
),
)
)
volume_mounts.append(
V1VolumeMount(
name=mount_name,
mount_path=mount.dst_path,
read_only=(
"w" not in mount.permissions and "m" not in mount.permissions
),
)
)
security_context.privileged = True
else:
raise TypeError(f"unknown mount type {mount}")
container = V1Container(
command=[role.entrypoint] + role.args,
image=role.image,
name=name,
env=[
(
V1EnvVar(
name=name,
value_from=V1EnvVarSource(
field_ref=V1ObjectFieldSelector(
field_path=value.removeprefix(PLACEHOLDER_FIELD_PATH)
)
),
)
if value.startswith(PLACEHOLDER_FIELD_PATH)
else V1EnvVar(
name=name,
value=value,
)
)
for name, value in role.env.items()
],
resources=resources,
ports=[
V1ContainerPort(
name=name,
container_port=port,
)
for name, port in role.port_map.items()
],
volume_mounts=volume_mounts,
security_context=security_context,
)
return V1Pod(
spec=V1PodSpec(
containers=[container],
restart_policy="Never",
service_account_name=service_account,
volumes=volumes,
node_selector=node_selector,
),
metadata=V1ObjectMeta(
annotations={
# Disable the istio sidecar as it prevents the containers from
# exiting once finished.
ANNOTATION_ISTIO_SIDECAR: "false",
},
labels={},
),
)
[docs]
def app_to_resource(
app: AppDef,
queue: str,
service_account: str | None,
priority_class: str | None = None,
reserved_millicpu: int = RESERVED_MILLICPU,
reserved_memmb: int = RESERVED_MEMMB,
efa_device_count: int | None = None,
) -> dict[str, Any]:
"""
app_to_resource creates a volcano job kubernetes resource definition from
the provided AppDef. The resource definition can be used to launch the
app on Kubernetes.
To support macros we generate one task per replica instead of using the
volcano `replicas` field since macros change the arguments on a per
replica basis.
Volcano has two levels of retries: one at the task level and one at the
job level. When using the APPLICATION retry policy, the job level retry
count is set to the minimum of the max_retries of the roles.
"""
tasks = []
unique_app_id = normalize_str(make_unique(app.name))
for role_idx, role in enumerate(app.roles):
for replica_id in range(role.num_replicas):
values = macros.Values(
img_root="",
app_id=unique_app_id,
replica_id=str(replica_id),
rank0_env=f"VC_{normalize_str(app.roles[0].name)}_0_HOSTS".upper(),
)
if role_idx == 0 and replica_id == 0:
values.rank0_env = "TORCHX_RANK0_HOST"
name = normalize_str(f"{role.name}-{replica_id}")
replica_role = values.apply(role)
if role_idx == 0 and replica_id == 0:
replica_role.env["TORCHX_RANK0_HOST"] = "localhost"
replica_role.env["TORCHX_IMAGE"] = replica_role.image
pod = role_to_pod(
name,
replica_role,
service_account,
reserved_millicpu,
reserved_memmb,
efa_device_count,
)
if pod_overlay := get_overlay(role, "kubernetes", "V1Pod"):
_apply_pod_overlay(pod, pod_overlay)
pod.metadata.labels.update(
pod_labels(
app=app,
role_idx=role_idx,
role=role,
replica_id=replica_id,
app_id=unique_app_id,
)
)
task: dict[str, Any] = {
"replicas": 1,
"name": name,
"template": pod,
}
if role.max_retries > 0:
task["maxRetry"] = role.max_retries
task["policies"] = RETRY_POLICIES[role.retry_policy]
msg = f"""
Role {role.name} configured with restarts: {role.max_retries}. As of 1.4.0 Volcano
does NOT support retries correctly. More info: https://github.com/volcano-sh/volcano/issues/1651
"""
warnings.warn(msg)
if role.min_replicas is not None:
# first min_replicas tasks are required, afterward optional
task["minAvailable"] = 1 if replica_id < role.min_replicas else 0
tasks.append(task)
job_retries = min(role.max_retries for role in app.roles)
job_spec = {
"schedulerName": "volcano",
"queue": queue,
"tasks": tasks,
"maxRetry": job_retries,
"plugins": {
# https://github.com/volcano-sh/volcano/issues/533
"svc": ["--publish-not-ready-addresses"],
"env": [],
},
}
if priority_class is not None:
job_spec["priorityClassName"] = priority_class
resource: dict[str, Any] = {
"apiVersion": "batch.volcano.sh/v1alpha1",
"kind": "Job",
"metadata": {"name": f"{unique_app_id}"},
"spec": job_spec,
}
return resource
[docs]
@dataclass
class KubernetesJob:
images_to_push: dict[str, tuple[str, str]]
resource: dict[str, Any]
def __str__(self) -> str:
return yaml.dump(sanitize_for_serialization(self.resource))
def __repr__(self) -> str:
return str(self)
@dataclass
class Opts(StructuredOpts):
"""Typed configuration options for KubernetesScheduler."""
queue: str
"""Volcano queue to schedule job in."""
namespace: str = "default"
"""Kubernetes namespace to schedule job in."""
service_account: str | None = None
"""The service account name to set on the pod specs."""
priority_class: str | None = None
"""The name of the PriorityClass to set on the job specs."""
validate_spec: bool = True
"""Validate job spec using Kubernetes API dry-run before submission."""
reserved_millicpu: int = RESERVED_MILLICPU
"""Amount of CPU in millicores to reserve for Kubernetes system overhead (default: 100)."""
reserved_memmb: int = RESERVED_MEMMB
"""Amount of memory in MB to reserve for Kubernetes system overhead (default: 1024)."""
image_repo: str | None = None
"""The image repository to use when pushing patched images, must have push access."""
efa_device_count: int | None = None
"""EFA device count override: None/unset=use resource spec, 0=remove EFA, N>0=set EFA count to N."""
KubernetesOpts = Opts
[docs]
class KubernetesScheduler(DockerWorkspaceMixin, Scheduler[Opts]):
"""
KubernetesScheduler is a TorchX scheduling interface to Kubernetes.
Important: Volcano is required to be installed on the Kubernetes cluster.
TorchX requires gang scheduling for multi-replica/multi-role execution
and Volcano is currently the only supported scheduler with Kubernetes.
For installation instructions see: https://github.com/volcano-sh/volcano
This has been confirmed to work with Volcano v1.3.0 and Kubernetes versions
v1.18-1.21. See https://github.com/meta-pytorch/torchx/issues/120 which is
tracking Volcano support for Kubernetes v1.22.
.. note::
AppDefs that have more than 0 retries may not be displayed as pods if they failed.
This occurs due to known bug in Volcano(as per 1.4.0 release):
https://github.com/volcano-sh/volcano/issues/1651
.. code-block:: bash
$ pip install torchx[kubernetes]
$ torchx run --scheduler kubernetes --scheduler_args namespace=default,queue=test utils.echo --image alpine:latest --msg hello
kubernetes://torchx_user/1234
$ torchx status kubernetes://torchx_user/1234
...
**Cancellation**
Canceling a job aborts it while preserving the job spec for inspection
and cloning via kubectl apply. Use the delete command to remove the job entirely:
.. code-block:: bash
$ torchx cancel kubernetes://namespace/jobname # abort, preserves spec
$ torchx delete kubernetes://namespace/jobname # delete completely
**Config Options**
.. runopts::
class: torchx.schedulers.kubernetes_scheduler.create_scheduler
**Mounts**
Mounting external filesystems/volumes is via the HostPath and
PersistentVolumeClaim support.
* hostPath volumes: ``type=bind,src=<host path>,dst=<container path>[,readonly]``
* PersistentVolumeClaim: ``type=volume,src=<claim>,dst=<container path>[,readonly]``
* host devices: ``type=device,src=/dev/foo[,dst=<container path>][,perm=rwm]``
If you specify a host device the job will run in privileged mode since
Kubernetes doesn't expose a way to pass `--device` to the underlying
container runtime. Users should prefer to use device plugins.
See :py:func:`torchx.specs.parse_mounts` for more info.
External docs: https://kubernetes.io/docs/concepts/storage/persistent-volumes/
**Resources / Allocation**
To select a specific machine type you can add a capability to your resources
with ``node.kubernetes.io/instance-type`` which will constrain the launched
jobs to nodes of that instance type.
>>> from torchx import specs
>>> specs.Resource(
... cpu=4,
... memMB=16000,
... gpu=2,
... capabilities={
... "node.kubernetes.io/instance-type": "<cloud instance type>",
... },
... )
Resource(...)
Kubernetes may reserve some memory for the host. TorchX assumes you're
scheduling on whole hosts and thus will automatically reduce the resource
request by a small amount to account for the node reserved CPU and memory.
If you run into scheduling issues you may need to reduce the requested CPU
and memory from the host values.
**Compatibility**
.. compatibility::
type: scheduler
features:
cancel: true
logs: true
distributed: true
describe: |
Partial support. KubernetesScheduler will return job and replica
status but does not provide the complete original AppSpec.
workspaces: true
mounts: true
elasticity: Requires Volcano >1.6
"""
def __init__(
self,
session_name: str,
client: "ApiClient | None" = None,
docker_client: "DockerClient | None" = None,
) -> None:
# NOTE: make sure any new init options are supported in create_scheduler(...)
super().__init__("kubernetes", session_name, docker_client=docker_client)
self._client = client
def _api_client(self) -> "ApiClient":
from kubernetes import client, config
c = self._client
if c is None:
configuration = client.Configuration()
try:
# Try in-cluster config first (for pods with ServiceAccount)
config.load_incluster_config(client_configuration=configuration)
except config.ConfigException:
# Fall back to kubeconfig (for local development)
try:
config.load_kube_config(client_configuration=configuration)
except config.ConfigException as e:
warnings.warn(f"failed to load kube config: {e}", stacklevel=2)
c = self._client = client.ApiClient(configuration)
return c
def _custom_objects_api(self) -> "CustomObjectsApi":
from kubernetes import client
return client.CustomObjectsApi(self._api_client())
def _get_job_name_from_exception(self, e: "ApiException") -> str | None:
try:
return json.loads(e.body)["details"]["name"]
except Exception:
logger.exception("unable to retrieve job name from api exception")
return None
def _get_active_context(self) -> dict[str, Any]:
from kubernetes import config
contexts, active_context = config.list_kube_config_contexts()
return active_context
[docs]
def schedule(self, dryrun_info: AppDryRunInfo[KubernetesJob]) -> str:
from kubernetes.client.rest import ApiException
namespace = dryrun_info.cfg.get("namespace") or "default"
images_to_push = dryrun_info.request.images_to_push
self.push_images(images_to_push)
resource = dryrun_info.request.resource
try:
resp = self._custom_objects_api().create_namespaced_custom_object(
group="batch.volcano.sh",
version="v1alpha1",
namespace=namespace,
plural="jobs",
body=resource,
)
except ApiException as e:
if e.status == 409 and e.reason == "Conflict":
job_name = self._get_job_name_from_exception(e)
raise ValueError(
f"Job `{job_name}` already exists. This seems like a transient exception, try resubmitting job"
) from e
else:
raise
return f"{namespace}:{resp['metadata']['name']}"
def _submit_dryrun(self, app: AppDef, cfg: Opts) -> AppDryRunInfo[KubernetesJob]:
queue = cfg.get("queue")
if not isinstance(queue, str):
raise TypeError(f"config value 'queue' must be a string, got {queue}")
# map any local images to the remote image
images_to_push = self.dryrun_push_images(app, cast(Mapping[str, CfgVal], cfg))
service_account = cfg.get("service_account")
assert service_account is None or isinstance(
service_account, str
), "service_account must be a str"
priority_class = cfg.get("priority_class")
assert priority_class is None or isinstance(
priority_class, str
), "priority_class must be a str"
reserved_millicpu = cfg.get("reserved_millicpu")
if reserved_millicpu is None:
reserved_millicpu = RESERVED_MILLICPU
assert isinstance(reserved_millicpu, int), "reserved_millicpu must be an int"
reserved_memmb = cfg.get("reserved_memmb")
if reserved_memmb is None:
reserved_memmb = RESERVED_MEMMB
assert isinstance(reserved_memmb, int), "reserved_memmb must be an int"
efa_device_count = cfg.get("efa_device_count")
assert efa_device_count is None or isinstance(
efa_device_count, int
), "efa_device_count must be an int or None"
resource = app_to_resource(
app,
queue,
service_account,
priority_class,
reserved_millicpu,
reserved_memmb,
efa_device_count,
)
if cfg.get("validate_spec"):
try:
self._custom_objects_api().create_namespaced_custom_object(
group="batch.volcano.sh",
version="v1alpha1",
namespace=cfg.get("namespace") or "default",
plural="jobs",
body=resource,
dry_run="All",
)
except Exception as e:
from kubernetes.client.rest import ApiException
if isinstance(e, ApiException):
raise ValueError(f"Invalid job spec: {e.reason}") from e
raise
job_name = resource["metadata"]["name"]
for task in resource["spec"]["tasks"]:
task_name = task["name"]
replicas = task.get("replicas", 1)
max_index = replicas - 1
pod_name = f"{job_name}-{task_name}-{max_index}"
if len(pod_name) > 63:
raise ValueError(
f"Pod name '{pod_name}' ({len(pod_name)} chars) exceeds 63 character limit. "
f"Shorten app.name or role names"
)
req = KubernetesJob(
resource=resource,
images_to_push=images_to_push,
)
return AppDryRunInfo(req, repr)
def _validate(self, app: AppDef, scheduler: str, cfg: Opts) -> None:
# Skip validation step
pass
def _cancel_existing(self, app_id: str) -> None:
"""
Abort a Volcano job while preserving the spec for inspection.
"""
namespace, name = app_id.split(":")
vcjob = self._custom_objects_api().get_namespaced_custom_object(
group="batch.volcano.sh",
version="v1alpha1",
namespace=namespace,
plural="jobs",
name=name,
)
vcjob.setdefault("status", {}).setdefault("state", {})["phase"] = "Aborted"
self._custom_objects_api().replace_namespaced_custom_object_status(
group="batch.volcano.sh",
version="v1alpha1",
namespace=namespace,
plural="jobs",
name=name,
body=vcjob,
)
def _delete_existing(self, app_id: str) -> None:
"""
Delete a Volcano job completely from the cluster.
"""
namespace, name = app_id.split(":")
self._custom_objects_api().delete_namespaced_custom_object(
group="batch.volcano.sh",
version="v1alpha1",
namespace=namespace,
plural="jobs",
name=name,
)
def _run_opts(self) -> runopts:
return Opts.as_runopts()
[docs]
def describe(self, app_id: str) -> DescribeAppResponse | None:
from kubernetes import client
from kubernetes.client.rest import ApiException
namespace, name = app_id.split(":")
roles = {}
roles_statuses = {}
try:
resp = self._custom_objects_api().get_namespaced_custom_object_status(
group="batch.volcano.sh",
version="v1alpha1",
namespace=namespace,
plural="jobs",
name=name,
)
except ApiException as e:
if e.status == 404:
return None
raise
status = resp.get("status")
if status:
state_str = status["state"]["phase"]
app_state = JOB_STATE[state_str]
TASK_STATUS_COUNT = "taskStatusCount"
if TASK_STATUS_COUNT in status:
for task_name, task_status in status[TASK_STATUS_COUNT].items():
role, _, idx = task_name.rpartition("-")
state_str = next(iter(task_status["phase"].keys()))
state = TASK_STATE[state_str]
if role not in roles:
roles[role] = Role(name=role, num_replicas=0, image="")
roles_statuses[role] = RoleStatus(role, [])
roles[role].num_replicas += 1
# Pod name follows the pattern: {job_name}-{task_name}-0
# Get the pod to retrieve its IP address
pod_name_k8s = f"{name}-{task_name}-0"
hostname = ""
try:
core_api = client.CoreV1Api(self._api_client())
pod = core_api.read_namespaced_pod(
name=pod_name_k8s, namespace=namespace
)
pod_ip = pod.status.pod_ip
if pod_ip is not None:
# Convert IP to dashed format (e.g., 10.244.1.5 -> 10-244-1-5)
pod_ip_dashed = pod_ip.replace(".", "-")
# Kubernetes DNS = <pod-ip-dashed>.<namespace>.pod.cluster.local
# Note: This will only be useful if the client using the IPs is in the cluster.
hostname = f"{pod_ip_dashed}.{namespace}.pod.cluster.local"
except ApiException:
# Pod not found - hostname remains empty
pass
roles_statuses[role].replicas.append(
ReplicaStatus(
id=int(idx), role=role, state=state, hostname=hostname
)
)
else:
app_state = AppState.UNKNOWN
return DescribeAppResponse(
app_id=app_id,
roles=list(roles.values()),
roles_statuses=list(roles_statuses.values()),
state=app_state,
)
[docs]
def describe_native(self, app_id: str) -> AppDryRunInfo[KubernetesJob] | None:
"""Reads the live Volcano job resource back from the cluster.
``images_to_push`` is empty on read-back: workspace images were
already pushed at submit time.
"""
from kubernetes.client.rest import ApiException
namespace, name = app_id.split(":")
try:
resource = self._custom_objects_api().get_namespaced_custom_object(
group="batch.volcano.sh",
version="v1alpha1",
namespace=namespace,
plural="jobs",
name=name,
)
except ApiException as e:
if e.status == 404:
return None
raise
return AppDryRunInfo(KubernetesJob(images_to_push={}, resource=resource), repr)
[docs]
def log_iter(
self,
app_id: str,
role_name: str,
k: int = 0,
regex: str | None = None,
since: datetime | None = None,
until: datetime | None = None,
should_tail: bool = False,
streams: Stream | None = None,
) -> Iterable[str]:
assert until is None, "kubernetes API doesn't support until"
if streams not in (None, Stream.COMBINED):
raise ValueError("KubernetesScheduler only supports COMBINED log stream")
from kubernetes import client, watch
namespace, name = app_id.split(":")
pod_name = normalize_str(f"{name}-{role_name}-{k}-0")
core_api = client.CoreV1Api(self._api_client())
pod = core_api.read_namespaced_pod(name=pod_name, namespace=namespace)
args: dict[str, object] = {
"name": pod_name,
"namespace": namespace,
"timestamps": True,
}
if since is not None:
args["since_seconds"] = (datetime.now() - since).total_seconds()
for container in pod.spec.containers:
args["container"] = container.name
if should_tail:
w = watch.Watch()
iterator = (
f"{line}\n"
for line in w.stream(core_api.read_namespaced_pod_log, **args)
)
else:
resp = core_api.read_namespaced_pod_log(**args)
iterator = split_lines(resp)
if regex:
iterator = filter_regex(regex, iterator)
for line in iterator:
yield f"{prefix_container_name(container.name, role_name, k)}{line}"
[docs]
def list(self, cfg: Mapping[str, CfgVal] | None = None) -> list[ListAppResponse]:
from kubernetes.config import ConfigException
namespace = (cfg or {}).get("namespace")
if not namespace:
try:
# in-cluster pods have no kubeconfig; the context may also
# omit `namespace`
namespace = self._get_active_context()["context"].get("namespace")
except ConfigException:
namespace = None
namespace = namespace or "default"
resp = self._custom_objects_api().list_namespaced_custom_object(
group="batch.volcano.sh",
version="v1alpha1",
namespace=namespace,
plural="jobs",
timeout_seconds=30,
)
return [
ListAppResponse(
app_id=f"{namespace}:{app['metadata']['name']}",
# just-created jobs may not have `status.state.phase` yet
state=JOB_STATE.get(
app.get("status", {}).get("state", {}).get("phase", ""),
AppState.UNKNOWN,
),
)
for app in resp["items"]
]
[docs]
def create_scheduler(
session_name: str,
client: "ApiClient | None" = None,
docker_client: "DockerClient | None" = None,
**kwargs: Any,
) -> KubernetesScheduler:
return KubernetesScheduler(
session_name=session_name,
client=client,
docker_client=docker_client,
)
[docs]
def pod_labels(
app: AppDef, role_idx: int, role: Role, replica_id: int, app_id: str
) -> dict[str, str]:
def clean(label_value: str) -> str:
# cleans the provided `label_value` to make it compliant
# to pod label specs as described in
# https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/
#
# Valid label value:
# must be 63 characters or less (can be empty),
# unless empty, must begin and end with an alphanumeric character ([a-z0-9A-Z]),
# could contain dashes (-), underscores (_), dots (.), and alphanumerics between.
# Replace invalid characters (allow: alphanum, -, _, .) with "."
label_value = re.sub(r"[^A-Za-z0-9\-_.]", ".", label_value)
# Replace leading non-alphanumeric with "."
label_value = re.sub(r"^[^A-Za-z0-9]+", ".", label_value)
# Replace trailing non-alphanumeric with "."
label_value = re.sub(r"[^A-Za-z0-9]+$", ".", label_value)
# Trim to 63 characters
return label_value[:63]
return {
LABEL_VERSION: clean(torchx.__version__),
LABEL_APP_NAME: clean(app.name),
LABEL_ROLE_INDEX: str(role_idx),
LABEL_ROLE_NAME: clean(role.name),
LABEL_REPLICA_ID: str(replica_id),
LABEL_KUBE_APP_NAME: clean(app.name),
LABEL_ORGANIZATION: "torchx.pytorch.org",
LABEL_UNIQUE_NAME: clean(app_id),
}
