Source code for torchx.schedulers.gcp_batch_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 GCP Batch scheduler which can be used to run TorchX
components directly on GCP Batch.

This scheduler is in prototype stage and may change without notice.


You need to have a GCP project configured to use Batch by enabling and setting it up.


from dataclasses import dataclass
from datetime import datetime
from typing import Any, Dict, Iterable, List, Optional, TYPE_CHECKING

import torchx
import yaml

from torchx.schedulers.api import (
from torchx.schedulers.ids import make_unique
from torchx.specs.api import AppDef, AppState, macros, Resource, Role, runopts
from torchx.util.strings import normalize_str
from typing_extensions import TypedDict

    from import batch_v1

JOB_STATE: Dict[str, AppState] = {
    "RUNNING": AppState.RUNNING,
    "FAILED": AppState.FAILED,

GPU_COUNT_TO_TYPE: Dict[int, str] = {
    1: "a2-highgpu-1g",
    2: "a2-highgpu-2g",
    4: "a2-highgpu-4g",
    8: "a2-highgpu-8g",
    16: "a2-highgpu-16g",

GPU_TYPE_TO_COUNT: Dict[str, int] = {v: k for k, v in GPU_COUNT_TO_TYPE.items()}

LABEL_VERSION: str = "torchx_version"
LABEL_APP_NAME: str = "torchx_app_name"

DEFAULT_LOC: str = "us-central1"

# TODO Remove LOCATIONS list once Batch supports all locations
# or when there is an API to query locations supported by Batch
LOCATIONS: List[str] = [

BATCH_LOGGER_NAME = "batch_task_logs"

[docs]@dataclass class GCPBatchJob: name: str project: str location: str job_def: "batch_v1.Job" def __str__(self) -> str: return yaml.dump(self.job_def) def __repr__(self) -> str: return str(self)
class GCPBatchOpts(TypedDict, total=False): project: Optional[str] location: Optional[str]
[docs]class GCPBatchScheduler(Scheduler[GCPBatchOpts]): """ GCPBatchScheduler is a TorchX scheduling interface to GCP Batch. .. code-block:: bash $ pip install torchx[gcp_batch] $ torchx run --scheduler gcp_batch utils.echo --msg hello # This launches a job with app handle like gcp_batch://torchx/project:location:app_id1234 and prints it $ torchx status gcp_batch://torchx/project:location:app_id1234 ... Authentication is loaded from the environment using the gcloud credential handling. **Config Options** .. runopts:: class: torchx.schedulers.gcp_batch_scheduler.create_scheduler **Compatibility** .. compatibility:: type: scheduler features: cancel: true logs: true describe: true distributed: true workspaces: false mounts: false elasticity: false """ def __init__( self, session_name: str, # pyre-fixme[2]: Parameter annotation cannot be `Any`. client: Optional[Any] = None, ) -> None: # NOTE: make sure any new init options are supported in create_scheduler(...) Scheduler.__init__(self, "gcp_batch", session_name) # pyre-fixme[4]: Attribute annotation cannot be `Any`. self.__client = client @property # pyre-fixme[3]: Return annotation cannot be `Any`. def _client(self) -> Any: from google.api_core import gapic_v1 from import batch_v1 c = self.__client if c is None: client_info = gapic_v1.client_info.ClientInfo( user_agent=f"TorchX/{torchx.__version__}" ) c = self.__client = batch_v1.BatchServiceClient(client_info=client_info) return c
[docs] def schedule(self, dryrun_info: AppDryRunInfo[GCPBatchJob]) -> str: from import batch_v1 req = dryrun_info.request assert req is not None, f"{dryrun_info} missing request" request = batch_v1.CreateJobRequest( parent=f"projects/{req.project}/locations/{req.location}", job=req.job_def,, ) response = self._client.create_job(request=request) return f"{req.project}:{req.location}:{}"
def _app_to_job(self, app: AppDef) -> "batch_v1.Job": from import batch_v1 name = normalize_str(make_unique( taskGroups = [] allocationPolicy = None # 1. Convert role to task # TODO implement retry_policy, mount conversion # NOTE: Supports only one role for now as GCP Batch supports only one TaskGroup # which is ok to start with as most components have only one role for role_idx, role in enumerate(app.roles): values = macros.Values( img_root="", app_id=name, replica_id=str(0), rank0_env=("BATCH_MAIN_NODE_HOSTNAME"), ) role_dict = values.apply(role) role_dict.env["TORCHX_ROLE_IDX"] = str(role_idx) role_dict.env["TORCHX_ROLE_NAME"] = str( resource = role_dict.resource res = batch_v1.ComputeResource() cpu = resource.cpu if cpu <= 0: cpu = 1 MILLI = 1000 res.cpu_milli = cpu * MILLI memMB = resource.memMB if memMB < 0: raise ValueError( f"memMB should to be set to a positive value, got {memMB}" ) res.memory_mib = memMB # TODO support named resources # Using v100 as default GPU type as a100 does not allow changing count for now # TODO See if there is a better default GPU type if resource.gpu > 0: if resource.gpu not in GPU_COUNT_TO_TYPE: raise ValueError( f"gpu should to be set to one of these values: {GPU_COUNT_TO_TYPE.keys()}" ) machineType = GPU_COUNT_TO_TYPE[resource.gpu] allocationPolicy = batch_v1.AllocationPolicy( instances=[ batch_v1.AllocationPolicy.InstancePolicyOrTemplate( install_gpu_drivers=True, policy=batch_v1.AllocationPolicy.InstancePolicy( machine_type=machineType, ), ) ], ) print(f"Using GPUs of type: {machineType}") # Configure host firewall rules to accept ingress communication config_network_runnable = batch_v1.Runnable( script=batch_v1.Runnable.Script( text="/sbin/iptables -A INPUT -j ACCEPT" ) ) runnable = batch_v1.Runnable( container=batch_v1.Runnable.Container( image_uri=role_dict.image, commands=[role_dict.entrypoint] + role_dict.args, entrypoint="", # Configure docker to use the host network stack to communicate with containers/other hosts in the same network options="--net host", ) ) ts = batch_v1.TaskSpec( runnables=[config_network_runnable, runnable], environment=batch_v1.Environment(variables=role_dict.env), max_retry_count=role_dict.max_retries, compute_resource=res, ) task_env = [ batch_v1.Environment(variables={"TORCHX_REPLICA_IDX": str(i)}) for i in range(role_dict.num_replicas) ] tg = batch_v1.TaskGroup( task_spec=ts, task_count=role_dict.num_replicas, task_count_per_node=1, task_environments=task_env, require_hosts_file=True, ) taskGroups.append(tg) # 2. Convert AppDef to Job job = batch_v1.Job( name=name, task_groups=taskGroups, allocation_policy=allocationPolicy, logs_policy=batch_v1.LogsPolicy( destination=batch_v1.LogsPolicy.Destination.CLOUD_LOGGING, ), # NOTE: GCP Batch does not allow label names with "." labels={ LABEL_VERSION: torchx.__version__.replace(".", "-"), LABEL_APP_NAME: name, }, ) return job def _get_project(self) -> str: from import runtimeconfig return runtimeconfig.Client().project def _submit_dryrun( self, app: AppDef, cfg: GCPBatchOpts ) -> AppDryRunInfo[GCPBatchJob]: proj = cfg.get("project") if proj is None: proj = self._get_project() assert proj is not None and isinstance(proj, str), "project must be a str" loc = cfg.get("location") assert loc is not None and isinstance(loc, str), "location must be a str" job = self._app_to_job(app) # Convert JobDef + BatchOpts to GCPBatchJob req = GCPBatchJob( name=str(, project=proj, location=loc, job_def=job, ) return AppDryRunInfo(req, repr)
[docs] def run_opts(self) -> runopts: opts = runopts() opts.add( "project", type_=str, help="Name of the GCP project. Defaults to the configured GCP project in the environment", ) opts.add( "location", type_=str, default=DEFAULT_LOC, help=f"Name of the location to schedule the job in. Defaults to {DEFAULT_LOC}", ) return opts
def _app_id_to_job_full_name(self, app_id: str) -> str: """ app_id format: f"{project}:{location}:{name}" job_full_name format: f"projects/{project}/locations/{location}/jobs/{name}" where 'name' was created uniquely for the job from the app name """ app_id_splits = app_id.split(":") if len(app_id_splits) != 3: raise ValueError(f"app_id not in expected format: {app_id}") return f"projects/{app_id_splits[0]}/locations/{app_id_splits[1]}/jobs/{app_id_splits[2]}" def _get_job(self, app_id: str) -> "batch_v1.Job": from import batch_v1 job_name = self._app_id_to_job_full_name(app_id) request = batch_v1.GetJobRequest( name=job_name, ) return self._client.get_job(request=request)
[docs] def describe(self, app_id: str) -> Optional[DescribeAppResponse]: job = self._get_job(app_id) if job is None: print(f"app not found: {app_id}") return None gpu = 0 if len(job.allocation_policy.instances) != 0: gpu_type = job.allocation_policy.instances[0].policy.machine_type gpu = GPU_TYPE_TO_COUNT[gpu_type] roles = {} for tg in job.task_groups: env = tg.task_spec.environment.variables role = env["TORCHX_ROLE_NAME"] container = tg.task_spec.runnables[1].container roles[role] = Role( name=role, num_replicas=tg.task_count, image=container.image_uri, entrypoint=container.commands[0], args=list(container.commands[1:]), resource=Resource( cpu=int(tg.task_spec.compute_resource.cpu_milli / 1000), memMB=tg.task_spec.compute_resource.memory_mib, gpu=gpu, ), env=dict(env), max_retries=tg.task_spec.max_retry_count, ) # Map job -> DescribeAppResponse # TODO map role/replica status desc = DescribeAppResponse( app_id=app_id, state=JOB_STATE[], roles=list(roles.values()), ) return desc
[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]: if streams not in (None, Stream.COMBINED): raise ValueError("GCPBatchScheduler only supports COMBINED log stream") job = self._get_job(app_id) if not job: raise ValueError(f"app not found: {app_id}") job_uid = job.uid filters = [ f"labels.job_uid={job_uid}", f"labels.task_id:{job_uid}-group0-{k}", ] if since is not None: filters.append(f'timestamp>="{str(since.isoformat())}"') else: # gcloud logger.list by default only returns logs in the last 24 hours # Since many ML jobs can run longer add timestamp filter to get all logs filters.append(f'timestamp>="{str(datetime.fromtimestamp(0).isoformat())}"') if until is not None: filters.append(f'timestamp<="{str(until.isoformat())}"') if regex is not None: filters.append(f'textPayload =~ "{regex}"') filter = " AND ".join(filters) return self._batch_log_iter(filter)
def _batch_log_iter(self, filter: str) -> Iterable[str]: from import logging logger = logging.Client().logger(BATCH_LOGGER_NAME) for entry in logger.list_entries(filter_=filter): yield entry.payload + "\n" def _job_full_name_to_app_id(self, job_full_name: str) -> str: """ job_full_name format: f"projects/{project}/locations/{location}/jobs/{name}" app_id format: f"{project}:{location}:{name}" where 'name' was created uniquely for the job from the app name """ job_name_splits = job_full_name.split("/") if len(job_name_splits) != 6: raise ValueError(f"job full name not in expected format: {job_full_name}") return f"{job_name_splits[1]}:{job_name_splits[3]}:{job_name_splits[5]}"
[docs] def list(self) -> List[ListAppResponse]: all_jobs = [] proj = self._get_project() for loc in LOCATIONS: jobs = self._client.list_jobs(parent=f"projects/{proj}/locations/{loc}") all_jobs += jobs all_jobs.sort(key=lambda job: job.create_time.timestamp(), reverse=True) return [ ListAppResponse( app_id=self._job_full_name_to_app_id(, state=JOB_STATE[], ) for job in all_jobs ]
def _validate(self, app: AppDef, scheduler: str) -> None: # Skip validation step pass def _cancel_existing(self, app_id: str) -> None: from import batch_v1 job_name = self._app_id_to_job_full_name(app_id) request = batch_v1.DeleteJobRequest( name=job_name, reason="Killed via TorchX", ) self._client.delete_job(request=request)
[docs]def create_scheduler( session_name: str, # pyre-fixme[2]: Parameter annotation cannot be `Any`. client: Optional[Any] = None, **kwargs: object, ) -> GCPBatchScheduler: return GCPBatchScheduler( session_name=session_name, client=client, )


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources