Shortcuts

Source code for torchx.schedulers.kubernetes_mcad_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_MCAD scheduler which can be used to run TorchX
components on a Kubernetes cluster via the Multi-Cluster-Application-Dispatcher (MCAD).

Prerequisites
==============

TorchX Kubernetes_MCAD scheduler depends on AppWrapper + MCAD.

Install MCAD: 
See deploying Multi-Cluster-Application-Dispatcher guide 
https://github.com/project-codeflare/multi-cluster-app-dispatcher/blob/main/doc/deploy/deployment.md

This implementation requires MCAD v1.34.1 or higher.

TorchX uses `torch.distributed.run <https://pytorch.org/docs/stable/elastic/run.html>`_ to run distributed training.

Learn more about running distributed trainers :py:mod:`torchx.components.dist`

"""

import json
import logging
import re

import warnings
from dataclasses import dataclass
from datetime import datetime
from typing import (
    Any,
    cast,
    Dict,
    Iterable,
    List,
    Mapping,
    Optional,
    Tuple,
    TYPE_CHECKING,
)

import torchx
import yaml
from torchx.schedulers.api import (
    AppDryRunInfo,
    DescribeAppResponse,
    filter_regex,
    ListAppResponse,
    Scheduler,
    split_lines,
    Stream,
)
from torchx.schedulers.ids import make_unique
from torchx.specs.api import (
    AppDef,
    AppState,
    BindMount,
    CfgVal,
    DeviceMount,
    macros,
    ReplicaState,
    ReplicaStatus,
    Resource,
    RetryPolicy,
    Role,
    RoleStatus,
    runopts,
    VolumeMount,
)

from torchx.workspace.docker_workspace import DockerWorkspaceMixin
from typing_extensions import TypedDict

if TYPE_CHECKING:
    from docker import DockerClient
    from kubernetes.client import ApiClient, CustomObjectsApi
    from kubernetes.client.models import (  # noqa: F401 imported but unused
        V1Container,
        V1ContainerPort,
        V1EnvVar,
        V1Pod,
        V1PodSpec,
        V1ResourceRequirements,
        V1Service,
    )
    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

RETRY_POLICIES: Mapping[str, Iterable[Mapping[str, str]]] = {
    RetryPolicy.REPLICA: [],
    RetryPolicy.APPLICATION: [
        {"event": "PodEvicted", "action": "RestartJob"},
        {"event": "PodFailed", "action": "RestartJob"},
    ],
}

# The AppWrapper Status - holistic view of pods/services
JOB_STATE: Dict[str, AppState] = {
    # Pending is the AppWrapper condition waiting for scheduling by MCAD
    "Pending": AppState.PENDING,
    # Running is the AppWrapper condition in Running.
    "Running": AppState.RUNNING,
    # Deleted is the AppWrapper condition where the torchX job is cancelled
    # by the user, the AppWrapper is deleted by the user, or error handling
    "Deleted": AppState.CANCELLED,
    # Failed is the job finishes unexpectedly
    "Failed": AppState.FAILED,
}

TASK_STATE: Dict[str, ReplicaState] = {
    # Pending dispatch means the AppWrapped task is not yet scheduled by MCAD
    "Pending dispatch": ReplicaState.PENDING,
    # Pending means the task is scheduled by MCAD
    "pending": ReplicaState.PENDING,
    # Running means a task is running on the host.
    "running": 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"


[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, unique_app_id: str, namespace: str, role: Role, service_account: Optional[str], image_secret: Optional[str], coscheduler_name: Optional[str], priority_class_name: Optional[str], network: Optional[str], ) -> "V1Pod": from kubernetes.client.models import ( # noqa: F811 redefinition of unused V1Container, V1ContainerPort, V1EmptyDirVolumeSource, V1EnvVar, V1HostPathVolumeSource, V1LocalObjectReference, 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) for device_name, device_limit in resource.devices.items(): limits[device_name] = str(device_limit) 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}") torchx_env_var = [ V1EnvVar( name=name, value=value, ) for name, value in role.env.items() ] my_env_var = [ V1EnvVar( name=f"TORCHX_MCAD_{cleanup_str(role.name)}_0_HOSTS".upper().replace( "-", "" ), value=f"{unique_app_id}-0.{unique_app_id}", ) ] container = V1Container( command=[role.entrypoint] + role.args, image=role.image, name=name, env=torchx_env_var + my_env_var, 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, ) # Get correct formatting for image secret imagesecret = V1LocalObjectReference(name=image_secret) metadata = V1ObjectMeta( name=name, annotations={ # Disable the istio sidecar as it prevents the containers from # exiting once finished. ANNOTATION_ISTIO_SIDECAR: "false", }, labels={}, namespace=namespace, ) if network is not None: metadata.annotations.update({"k8s.v1.cni.cncf.io/networks": network}) return V1Pod( api_version="v1", kind="Pod", spec=V1PodSpec( containers=[container], hostname=name, subdomain=unique_app_id, image_pull_secrets=[imagesecret], restart_policy="Never", service_account_name=service_account, volumes=volumes, node_selector=node_selector, scheduler_name=coscheduler_name, priority_class_name=priority_class_name, ), metadata=metadata, )
def create_pod_group( app: AppDef, role: Role, role_idx: int, namespace: str, app_id: str ) -> "Dict[str, Any]": pod_group_name = app_id + "-pg" + str(role_idx) labels = object_labels(app, app_id) labels.update({"appwrapper.workload.codeflare.dev": app_id}) pod_group: Dict[str, Any] = { "apiVersion": "scheduling.sigs.k8s.io/v1alpha1", "kind": "PodGroup", "metadata": { "name": pod_group_name, "namespace": namespace, "labels": labels, }, "spec": { "minMember": role.num_replicas, }, } genericitem_pod_group: Dict[str, Any] = { "replicas": 1, "generictemplate": pod_group, } return genericitem_pod_group
[docs]def mcad_svc( app: AppDef, svc_name: str, namespace: str, service_port: str ) -> "V1Service": from kubernetes.client.models import ( # noqa: F401, F811 V1Container, V1ContainerPort, V1EmptyDirVolumeSource, V1EnvVar, V1HostPathVolumeSource, V1ObjectMeta, V1PersistentVolumeClaimVolumeSource, V1Pod, V1PodSpec, V1ResourceRequirements, V1SecurityContext, V1Service, V1ServicePort, V1ServiceSpec, V1ServiceStatus, V1Volume, V1VolumeMount, ) labels = object_labels(app, svc_name) return V1Service( api_version="v1", kind="Service", metadata=V1ObjectMeta( name=svc_name, namespace=namespace, labels=labels, ), spec=V1ServiceSpec( cluster_ip="None", publish_not_ready_addresses=True, ports=[ V1ServicePort( protocol="TCP", port=int(service_port), target_port=int(service_port), ) ], selector={LABEL_UNIQUE_NAME: svc_name}, session_affinity="None", type="ClusterIP", ), status=V1ServiceStatus( load_balancer={}, ), )
def cleanup_str(data: str) -> str: """ Invokes ``lower`` on thes string and removes all characters that do not satisfy ``[a-z0-9]`` pattern. This method is mostly used to make sure kubernetes scheduler gets the job name that does not violate its validation. """ if data.startswith("-"): data = data[1:] pattern = r"[a-z0-9\-]" return "".join(re.findall(pattern, data.lower())).lstrip("0123456789") def get_unique_truncated_appid(app: AppDef) -> str: """ Some Kubernetes objects need to have names that are 63 characters or less. When creating the unique app_id, this function calculates the max size to pass to make_unique. The PodGroup name includes 3 characters plus the role_id characters. The minimum number of characters for the unique identifier is 4. These amounts are taken into account. """ default_size = 14 uid_chars = 4 pg_chars = 3 + len(app.roles) size = 63 - (len(app.name) + uid_chars + pg_chars) unique_id_size = default_size if size > default_size else size if unique_id_size <= 3: msg = "Name size has too many characters for some Kubernetes objects. Truncating \ application name." warnings.warn(msg) end = 63 - uid_chars - pg_chars substring = app.name[0:end] app.name = substring unique_id_size = 3 unique_app_id = cleanup_str(make_unique(app.name, unique_id_size)) return unique_app_id
[docs]def get_port_for_service(app: AppDef) -> str: # Initialize port to default port = "29500" for role_idx, role in enumerate(app.roles): if role.port_map is None: continue for value in role.port_map.values(): port = str(value) if not (0 < int(port) <= 65535): msg = """Warning: port_map set to invalid port number. Value must be between 1-65535, with torchx default = 29500. Setting port to default = 29500""" port = "29500" warnings.warn(msg) return port
def enable_retry( job_spec: Dict[str, Any], appwrapper_retries: int, total_pods: int ) -> None: requeue_dict = { "timeInSeconds": 300, "maxTimeInSeconds": 0, "growthType": "exponential", "maxNumRequeuings": appwrapper_retries, } nested_specs = {"minAvailable": total_pods, "requeuing": requeue_dict} job_spec["schedulingSpec"] = nested_specs
[docs]def app_to_resource( app: AppDef, namespace: str, service_account: Optional[str], image_secret: Optional[str], coscheduler_name: Optional[str], priority_class_name: Optional[str], network: Optional[str], priority: Optional[int] = None, ) -> Dict[str, Any]: """ app_to_resource creates a AppWrapper/MCAD Kubernetes resource definition from the provided AppDef. The resource definition can be used to launch the app on Kubernetes. MCAD supports retries at the APPLICATION level. In the case of multiple TorchX Roles, the AppWrapper maximum number of retries count is set to the minimum of the max_retries of the roles. """ genericitems = [] unique_app_id = get_unique_truncated_appid(app) if coscheduler_name is not None: for role_idx, role in enumerate(app.roles): genericitem_pod_group = create_pod_group( app, role, role_idx, namespace, unique_app_id ) genericitems.append(genericitem_pod_group) 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"TORCHX_MCAD_{cleanup_str(app.roles[0].name)}_0_HOSTS".upper().replace( "-", "" ), ) if role_idx == 0 and replica_id == 0: values.rank0_env = "TORCHX_RANK0_HOST" name = cleanup_str(f"{unique_app_id}-{replica_id}") replica_role = values.apply(role) if role_idx == 0 and replica_id == 0: replica_role.env["TORCHX_RANK0_HOST"] = "localhost" pod = role_to_pod( name, unique_app_id, namespace, replica_role, service_account, image_secret, coscheduler_name, priority_class_name, network, ) pod.metadata.labels.update( pod_labels( app=app, role_idx=role_idx, role=role, replica_id=replica_id, coscheduler_name=coscheduler_name, app_id=unique_app_id, ) ) genericitem: Dict[str, Any] = { "replicas": 1, "generictemplate": pod, } genericitems.append(genericitem) """ Create Service: The selector will have the key 'appwrapper.workload.codeflare.dev', and the value will be the appwrapper name """ service_port = get_port_for_service(app) svc_obj = mcad_svc( app=app, svc_name=unique_app_id, namespace=namespace, service_port=service_port ) genericitem_svc: Dict[str, Any] = { "replicas": 1, "generictemplate": svc_obj, } genericitems.append(genericitem_svc) job_spec: Dict[str, Any] = { "resources": { "GenericItems": genericitems, }, } if priority is not None: job_spec["priority"] = priority appwrapper_retries = min(role.max_retries for role in app.roles) if appwrapper_retries > 0: total_pods = sum(role.num_replicas for role in app.roles) enable_retry(job_spec, appwrapper_retries, total_pods) resource: Dict[str, object] = { "apiVersion": "workload.codeflare.dev/v1beta1", "kind": "AppWrapper", "metadata": {"name": unique_app_id, "namespace": namespace}, "spec": job_spec, } return resource
# Helper function for MCAD generic items information -> TorchX Role
[docs]def get_role_information(generic_items: Iterable[Dict[str, Any]]) -> Dict[str, Any]: # Store unique role information roles = {} # nested dictionary keys # meta data information GT_KEY = "generictemplate" METADATA_KEY = "metadata" LABEL_KEY = "labels" ROLE_KEY = "torchx.pytorch.org/role-name" # containers information SPEC_KEY = "spec" CONTAINER_KEY = "containers" IMAGE_KEY = "image" ARGS_KEY = "command" ENV_KEY = "env" RESOURCE_KEY = "resources" PORTS_KEY = "ports" MOUNTS_KEY = "volumeMounts" # resource keys CPU_KEY = "cpu" GPU_KEY = "gpu" REQUEST_KEY = "requests" MEM_KEY = "memory" for generic_item in generic_items: if GT_KEY not in generic_item.keys(): continue gt_result = generic_item[GT_KEY] if METADATA_KEY not in gt_result.keys(): continue # Note: options in meta data : annotations, labels, name, namespace metadata_result = gt_result[METADATA_KEY] if LABEL_KEY not in metadata_result.keys(): continue label_result = metadata_result[LABEL_KEY] if ROLE_KEY not in label_result.keys(): continue role_name: str = label_result[ROLE_KEY] # save role if role_name not in roles: roles[role_name] = Role(name=role_name, num_replicas=0, image="") roles[role_name].num_replicas += 1 # Only get specs for first instance of TorchX role if SPEC_KEY not in gt_result.keys(): continue # Note: options in spec data: containers, hostname, imagePullSecrets, nodeSelector # restartPolicy, subdomain, volumes spec_result = gt_result[SPEC_KEY] if CONTAINER_KEY not in spec_result.keys(): continue container_result = spec_result[CONTAINER_KEY] if IMAGE_KEY not in container_result[0].keys(): continue roles[role_name].image = container_result[0][IMAGE_KEY] if ARGS_KEY not in container_result[0].keys(): continue roles[role_name].args = container_result[0][ARGS_KEY] if ENV_KEY not in container_result[0].keys(): continue roles[role_name].env = container_result[0][ENV_KEY] if RESOURCE_KEY not in container_result[0].keys(): continue roles[role_name].resources = container_result[0][RESOURCE_KEY] resource_req = Resource(cpu=-1, gpu=-1, memMB=-1) if CPU_KEY not in container_result[0][RESOURCE_KEY][REQUEST_KEY]: continue resource_req.cpu = container_result[0][RESOURCE_KEY][REQUEST_KEY][CPU_KEY] # Substring matching to accomodate different gpu types gpu_key_values = dict( filter( lambda item: GPU_KEY in item[0], container_result[0][RESOURCE_KEY][REQUEST_KEY].items(), ) ) if len(gpu_key_values) != 0: for key, value in gpu_key_values.items(): resource_req.gpu = value if MEM_KEY not in container_result[0][RESOURCE_KEY][REQUEST_KEY]: continue resource_req.memMB = container_result[0][RESOURCE_KEY][REQUEST_KEY][MEM_KEY] roles[role_name].resource = resource_req if PORTS_KEY not in container_result[0].keys(): continue roles[role_name].port_map = container_result[0][PORTS_KEY] if MOUNTS_KEY not in container_result[0].keys(): continue roles[role_name].mounts = container_result[0][MOUNTS_KEY] else: roles[role_name].num_replicas += 1 return roles
[docs]def get_appwrapper_status(app: Dict[str, str]) -> AppState: if "status" in app.keys(): # pyre-fixme return JOB_STATE[app["status"]["state"]] else: # Handle case where appwrapper is created but pending dispatch return JOB_STATE["Pending"]
# Does not handle not ready to dispatch case
[docs]def get_tasks_status_description(status: Dict[str, str]) -> Dict[str, int]: results = {} # Keys related to tasks and status KEY_RUN = "running" KEY_PEND = "pending" KEY_FAIL = "failed" KEY_SUCCESS = "Succeeded" if KEY_RUN in status.keys(): results[KEY_RUN] = status[KEY_RUN] if KEY_PEND in status.keys(): results[KEY_PEND] = status[KEY_PEND] if KEY_FAIL in status.keys(): results[KEY_FAIL] = status[KEY_FAIL] if KEY_SUCCESS in status.keys(): results[KEY_SUCCESS] = status[KEY_SUCCESS] return results
[docs]@dataclass class KubernetesMCADJob: images_to_push: Dict[str, Tuple[str, str]] resource: Dict[str, object] def __str__(self) -> str: return yaml.dump(sanitize_for_serialization(self.resource)) def __repr__(self) -> str: return str(self)
class KubernetesMCADOpts(TypedDict, total=False): namespace: Optional[str] image_repo: Optional[str] service_account: Optional[str] priority: Optional[int] priority_class_name: Optional[str] image_secret: Optional[str] coscheduler_name: Optional[str] network: Optional[str]
[docs]class KubernetesMCADScheduler(DockerWorkspaceMixin, Scheduler[KubernetesMCADOpts]): """ KubernetesMCADScheduler is a TorchX scheduling interface to Kubernetes. Important: AppWrapper/MCAD is required to be installed on the Kubernetes cluster. TorchX requires gang scheduling for multi-replica/multi-role execution. Note that AppWrapper/MCAD supports gang scheduling among any app-wrapped jobs on Kubernetes. However, for true gang scheduling AppWrapper/MCAD needs to be used with an additional Kubernetes co-scheduler. For installation instructions see: https://github.com/project-codeflare/multi-cluster-app-dispatcher/blob/main/doc/deploy/deployment.md This has been confirmed to work with MCAD main branch v1.34.1 or higher and OpenShift Kubernetes Client Version: 4.10.13 Server Version: 4.9.18 Kubernetes Version: v1.22.3+e790d7f .. code-block:: bash $ torchx run --scheduler kubernetes_mcad --scheduler_args namespace=default,image_repo=<your_image_repo> utils.echo --image alpine:latest --msg hello ... The TorchX-MCAD scheduler can be used with a secondary scheduler on Kubernetes. To enable this, the user must provide the name of the coscheduler. With this feature, a PodGroup is defined for each TorchX role and the coscheduler handles secondary scheduling on the Kubernetes cluster. For additional resources, see: 1. PodGroups and Coscheduling: https://github.com/kubernetes-sigs/scheduler-plugins/tree/release-1.24/pkg/coscheduling 2. Installing Secondary schedulers: https://github.com/kubernetes-sigs/scheduler-plugins/blob/release-1.24/doc/install.md 3. PodGroup CRD: https://github.com/kubernetes-sigs/scheduler-plugins/blob/release-1.24/config/crd/bases/scheduling.sigs.k8s.io_podgroups.yaml The MCAD scheduler supports priorities at the AppWrapper level and optionally at the pod level on clusters with PriorityClass definitions. At the AppWrapper level, higher integer values means higher priorities. Kubernetes clusters may have additional priorityClass definitions that can be applied at the pod level. While these different levels of priorities can be set independently, it is recommended to check with your Kubernetes cluster admin to see if additional guidance is in place. For more on Kubernetes PriorityClass, see: https://kubernetes.io/docs/concepts/scheduling-eviction/pod-priority-preemption/ . In order to use the network option, the Kubernetes cluster must have multus installed. For multus installation instructions and how to set up a network custom network attachment definition, see: https://github.com/k8snetworkplumbingwg/multus-cni/blob/master/docs/how-to-use.md **Config Options** .. runopts:: class: torchx.schedulers.kubernetes_mcad_scheduler.KubernetesMCADScheduler **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: true workspaces: true mounts: true elasticity: false """ def __init__( self, session_name: str, client: Optional["ApiClient"] = None, docker_client: Optional["DockerClient"] = None, ) -> None: # NOTE: make sure any new init options are supported in create_scheduler(...) super().__init__("kubernetes_mcad", 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: config.load_kube_config(client_configuration=configuration) except config.ConfigException as e: warnings.warn(f"failed to load kube config: {e}") c = self._client = client.ApiClient(configuration) return c def _custom_objects_api(self) -> "CustomObjectsApi": from kubernetes import client api_client = client.CustomObjectsApi(self._api_client()) return api_client def _get_job_name_from_exception(self, e: "ApiException") -> Optional[str]: try: return json.loads(e.body)["details"]["name"] except Exception as e: logger.exception("Unable to retrieve job name, got exception", e) 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[KubernetesMCADJob]) -> str: from kubernetes.client.rest import ApiException cfg = dryrun_info._cfg assert cfg is not None, f"{dryrun_info} missing cfg" namespace = 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="workload.codeflare.dev", version="v1beta1", namespace=namespace, plural="appwrappers", 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: KubernetesMCADOpts ) -> AppDryRunInfo[KubernetesMCADJob]: # map any local images to the remote image # images_to_push = self._update_app_images(app, cfg.get("image_repo")) 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 = cfg.get("priority") assert priority is None or isinstance(priority, int), "priority must be a int" image_secret = cfg.get("image_secret") assert image_secret is None or isinstance( image_secret, str ), "image_secret must be a str" if image_secret is not None and service_account is not None: msg = """Service Account and Image Secret names are both provided. Depending on the Service Account configuration, an ImagePullSecret may be defined in your Service Account. If this is the case, check service account and image secret configurations to understand the expected behavior for patched image push access.""" warnings.warn(msg) namespace = cfg.get("namespace") assert isinstance(namespace, str), "namespace must be a str" coscheduler_name = cfg.get("coscheduler_name") assert coscheduler_name is None or isinstance( coscheduler_name, str ), "coscheduler_name must be a string" priority_class_name = cfg.get("priority_class_name") assert priority_class_name is None or isinstance( priority_class_name, str ), "priority_class_name must be a string" network = cfg.get("network") assert network is None or isinstance(network, str), "network must be a string" resource = app_to_resource( app=app, namespace=namespace, service_account=service_account, image_secret=image_secret, coscheduler_name=coscheduler_name, priority_class_name=priority_class_name, network=network, priority=priority, ) req = KubernetesMCADJob( resource=resource, images_to_push=images_to_push, ) info = AppDryRunInfo(req, repr) info._app = app # pyre-fixme info._cfg = cfg return info def _validate(self, app: AppDef, scheduler: str) -> None: # Skip validation step pass def _cancel_existing(self, app_id: str) -> None: namespace, name = app_id.split(":") self._custom_objects_api().delete_namespaced_custom_object( group="workload.codeflare.dev", version="v1beta1", namespace=namespace, plural="appwrappers", name=name, )
[docs] def run_opts(self) -> runopts: opts = runopts() opts.add( "namespace", type_=str, help="Kubernetes namespace to schedule job in", default="default", ) opts.add( "image_repo", type_=str, help="The image repository to use when pushing patched images, must have push access. Ex: example.com/your/container", ) opts.add( "service_account", type_=str, help="The service account name to set on the pod specs", ) opts.add( "priority", type_=int, help="The priority level to set on the job specs. Higher integer value means higher priority", ) opts.add( "priority_class_name", type_=str, help="Pod specific priority level. Check with your Kubernetes cluster admin if Priority classes are defined on your system", ) opts.add( "image_secret", type_=str, help="The name of the Kubernetes/OpenShift secret set up for private images", ) opts.add( "coscheduler_name", type_=str, help="Option to run TorchX-MCAD with a co-scheduler. User must provide the co-scheduler name.", ) opts.add( "network", type_=str, help="Name of additional pod-to-pod network beyond default Kubernetes network", ) return opts
[docs] def describe(self, app_id: str) -> Optional[DescribeAppResponse]: namespace, name = app_id.split(":") from kubernetes.client.rest import ApiException roles = {} roles_statuses = {} # Production section api_instance = self._custom_objects_api group = "workload.codeflare.dev" version = "v1beta1" plural = "appwrappers" try: api_resp = api_instance().get_namespaced_custom_object( group, version, namespace, plural, name ) except ApiException as e: api_resp = {} if e.status == 404 and e.reason == "Not Found": raise ValueError( "Kubernetes client not found. Check access to Kubernetes cluster." ) from e elif e.status == 401 and e.reason == "Unauthorized": raise ValueError("Unauthorized Kubernetes access error.") from e else: raise task_status = [] if "status" in api_resp.keys(): status = api_resp["status"] tasks_results = get_tasks_status_description(status) # Handle case where waiting for dispatch if not tasks_results: tasks_results["Pending dispatch"] = ( len(api_resp["spec"]["resources"]["GenericItems"]) - 1 ) # Convert MCAD status to TorchX replica set status format # Warning: Status is not necessarily the match for a particular Replica ID for key, value in tasks_results.items(): for id in range(0, value): task_status.append(key) state = status["state"] app_state = JOB_STATE[state] # Roles spec = api_resp["spec"] resources = spec["resources"] generic_items = resources["GenericItems"] # Note MCAD service is not considered a TorchX role roles = get_role_information(generic_items) task_count = 0 for role in roles: msg = "Warning - MCAD does not report individual replica statuses, but overall task status. Replica id may not match status" warnings.warn(msg) roles_statuses[role] = RoleStatus(role, []) for idx in range(0, roles[role].num_replicas): state = TASK_STATE[task_status[task_count]] roles_statuses[role].replicas.append( ReplicaStatus(id=int(idx), role=role, state=state, hostname="") ) task_count += 1 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 log_iter( self, app_id: str, role_name: str, k: int = 0, regex: Optional[str] = None, since: Optional[datetime] = None, until: Optional[datetime] = None, should_tail: bool = False, streams: Optional[Stream] = None, ) -> Iterable[str]: assert until is None, "kubernetes API doesn't support until" if streams not in (None, Stream.COMBINED): raise ValueError( "KubernetesMCADScheduler only supports COMBINED log stream" ) from kubernetes import client, watch namespace, name = app_id.split(":") pod_name = cleanup_str(f"{name}-{k}") args: Dict[str, object] = { "name": pod_name, "namespace": namespace, "timestamps": True, } if since is not None: args["since_seconds"] = (datetime.now() - since).total_seconds() core_api = client.CoreV1Api(self._api_client()) if should_tail: w = watch.Watch() iterator = w.stream(core_api.read_namespaced_pod_log, **args) else: resp = core_api.read_namespaced_pod_log(**args) iterator = split_lines(resp) if regex: return filter_regex(regex, iterator) else: return iterator
[docs] def list(self) -> List[ListAppResponse]: active_context = self._get_active_context() namespace = active_context["context"]["namespace"] resp = self._custom_objects_api().list_namespaced_custom_object( group="workload.codeflare.dev", version="v1beta1", namespace=namespace, plural="appwrappers", timeout_seconds=30, ) return [ ListAppResponse( app_id=f"{namespace}:{item['metadata']['name']}", state=get_appwrapper_status(item), ) for item in resp["items"] ]
[docs]def create_scheduler( session_name: str, client: Optional["ApiClient"] = None, docker_client: Optional["DockerClient"] = None, **kwargs: Any, ) -> KubernetesMCADScheduler: return KubernetesMCADScheduler( session_name=session_name, client=client, docker_client=docker_client, )
def object_labels( app: AppDef, app_id: str, ) -> Dict[str, str]: return { LABEL_KUBE_APP_NAME: app.name, LABEL_ORGANIZATION: "torchx.pytorch.org", LABEL_UNIQUE_NAME: app_id, }
[docs]def pod_labels( app: AppDef, role_idx: int, role: Role, replica_id: int, coscheduler_name: Optional[str], app_id: str, ) -> Dict[str, str]: labels = object_labels(app, app_id) pod_labels = { LABEL_VERSION: torchx.__version__, LABEL_APP_NAME: app.name, LABEL_ROLE_INDEX: str(role_idx), LABEL_ROLE_NAME: role.name, LABEL_REPLICA_ID: str(replica_id), } if coscheduler_name is not None: pod_group = app_id + "-pg" + str(role_idx) pod_labels.update({"pod-group.scheduling.sigs.k8s.io": pod_group}) labels.update(pod_labels) return labels

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources