python-batch/batch.py

#!/usr/bin/env python3
#  Copyright 2023 Google LLC
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""
Tools to run Google Cloud Batch API
"""

__author__ = "J Ross Thomson drj@"
__version__ = "0.1.0"

import json
import os
import uuid
import yaml
import sys

from absl import app
from absl import flags
from google.api_core.operation import Operation
from google.cloud import batch_v1
from google.cloud import pubsub_v1

from typing import Iterable
from yaml.loader import SafeLoader

"""
We define multiple command line flags:
"""

FLAGS = flags.FLAGS

flags.DEFINE_string("config_file", None, "Config file in YAML")
flags.DEFINE_boolean("pubsub", False , "Run Pubsub Topic and Subscriber")
flags.DEFINE_string("previous_job_id", None, "For Pubsub restart, specifies topic to read from")
flags.DEFINE_string("project_id", None, "Google Cloud Project ID, not name")
flags.DEFINE_spaceseplist("volumes", None, "List of GCS paths to mount. Example, \"bucket_name1:mountpath1 bucket_name2:mountpath2\"" )
flags.DEFINE_boolean("create_job", False, "Creates job, otherwise just prints config.")
flags.DEFINE_boolean("list_jobs", False, "If true, list jobs for config.")
flags.DEFINE_string("delete_job", "", "Job name to delete.")
flags.DEFINE_boolean("debug", False, "If true, print debug info.")

# Required flag.
flags.mark_flag_as_required("config_file")

class PubSub:
  def __init__(self, job_id: str, config, previous_job_id=None) -> None:
      """Class to create Pub/Sub related cloud resources.

      Used to create a topic and a subscription. 
      Topic name identical to job_id and subscription is `sub-` + job_id.

      Args: 
          config: 
            Contains the structure defined by the Config.yaml file. 
          job_id: 
            The name of the job, which becomes the name of the pubsub topic.

      Raises:
        None

      `project_id` is set  based on config, env, or argv in that order.
      """
      self.config = config
      self.job_id = job_id

      if self.config["project_id"]: 
        self.project_id = self.config["project_id"] 
      if os.environ.get('GOOGLE_CLOUD_PROJECT'): 
        self.project_id  = os.environ.get('GOOGLE_CLOUD_PROJECT') 
      if FLAGS.project_id:
        self.project_id = FLAGS.project_id

      """
      Sending messages to the same region ensures they are received in order
      even when multiple publishers are used.
      """
      publisher_options = pubsub_v1.types.PublisherOptions(enable_message_ordering=True)
      client_options = {"api_endpoint": f'{self.config["region"]}-pubsub.googleapis.com:443'}
      self.publisher = pubsub_v1.PublisherClient(
          publisher_options=publisher_options, client_options=client_options
      )

  def create_topic(self):
    """
    Creates a topic with name the same as Job.

    The `topic_path` method creates a fully qualified identifier
    in the form `projects/{project_id}/topics/{topic_id}`

    Args:
      None

    Returns:
      None
    """
    self.topic_path = self.publisher.topic_path(self.project_id, self.job_id)
    self.topic = self.publisher.create_topic(request={"name": self.topic_path})
    print(f"Created topic {self.topic_path}\n", file=sys.stderr)

  def create_subscription(self):
    """
    Creates a subscription with name the same as `sub-` + Job.

    Subscription will be ordered and have exactly one time delivery.

    Args:

    Returns:
      None
    """

    self.subscriber = pubsub_v1.SubscriberClient()
    self.subscription_id = "sub-" + self.job_id

    self.topic_path = self.publisher.topic_path(self.project_id, self.job_id)
    self.subscription_path = self.subscriber.subscription_path(self.project_id, self.subscription_id)

    with self.subscriber:
      self.subscription = self.subscriber.create_subscription(
        request={
          "name": self.subscription_path, 
          "topic": self.topic_path,
          "enable_message_ordering": True,
          "enable_exactly_once_delivery": True,
        }
      )
    print(f"Created subscription {self.subscription_path}\n", file=sys.stderr)
    

  def publish_fifo_ids(self):
    """
    Publish integer ID messages to the pubsub topic.

    Simple queue create via the Pubsub Topic to create queue. 
    If the queue is not fully pulled, the project can restart with the same queue.
    Picking up where it left off.

    Args:
      None

    Returns:
      None
    """
    self.order_key = "fifo"
    for i in range(0, self.config["task_count"] ):
        # Data must be a bytestring
        data = str(i).encode("utf-8")
        # When you publish a message, the client returns a future.
        future = self.publisher.publish(self.topic_path, data=data, ordering_key=self.order_key)
        if FLAGS.debug:
          print(f'Future: {future.result()}, Message: {data}', file=sys.stderr)


    

class CloudBatchJobs:

  def __init__(self, job_id: str, config, env_vars) -> None:
    """
    Class to create all cloud resources for Batch Jobs

    Args:
      job_id: an arbitrary string to name the job.
      config: data structure from YAML to represent everything in the config file.

    Returns:
      None
    """
    self.config = config
    self.job_id = job_id
    self.env_vars = env_vars
    
    #set  "project_id"  based on config, env, or argv in that order.
  
    if self.config["project_id"]: 
      self.project_id = self.config["project_id"] 
    if os.environ.get('GOOGLE_CLOUD_PROJECT'): 
      self.project_id  = os.environ.get('GOOGLE_CLOUD_PROJECT') 
    if FLAGS.project_id:
      self.project_id = FLAGS.project_id

    self.client = batch_v1.BatchServiceClient()

  def _create_runnable(self) -> batch_v1.Runnable:

    self.runnable = batch_v1.Runnable()
    self.runnable.environment = batch_v1.Environment()
    self.runnable.environment.variables = self.env_vars

    if "container" in self.config:
      self.runnable.container = batch_v1.Runnable.Container()
      self.runnable.container.image_uri = self.config["container"]["image_uri"] 
      self.runnable.container.entrypoint =  self.config["container"]["entry_point"] 
      self.runnable.container.commands =  self.config["container"]["commands"] 
      self.runnable.container.options = "--privileged"
      if "install_gpu_drivers" in self.config:
        self.runnable.container.volumes.append("/var/lib/nvidia/lib64:/usr/local/nvidia/lib64")
        self.runnable.container.volumes.append("/var/lib/nvidia/bin:/usr/local/nvidia/bin")

    else:
      self.runnable.script = batch_v1.Runnable.Script()
      self.runnable.script.text = self.config["script_text"]
    return(self.runnable)

    
  def _create_task(self) -> batch_v1.TaskSpec:
    """
    _summary_

    :return: _description_
    :rtype: batch_v1.TaskSpec
    """
    self.task = batch_v1.TaskSpec()
    self.task.max_retry_count = 2
    self.task.max_run_duration = "604800s"

    self.task.volumes = []

    if "nfs_server" in self.config:

      nfs_server = batch_v1.NFS()
      nfs_server.server = self.config["nfs_server"]
      nfs_server.remote_path = self.config["nfs_remote_path"] if "nfs_remote_path" in self.config else "/nfshare"

      nfs_volume = batch_v1.Volume()
      nfs_volume.nfs = nfs_server
      nfs_volume.mount_path = self.config["nfs_path"] if "nfs_path" in self.config else "/mnt/nfs"

      self.task.volumes.append(nfs_volume)


# Use commmand line flags first, then config file, to set
# gcs bucket mounts
    if(FLAGS.volumes):
      self.volumes_list = []
      for volume_pair in FLAGS.volumes:
        (bucket_name, gcs_path) = volume_pair.split(":")
        self.volumes_list.append({"bucket_name":bucket_name, "gcs_path":gcs_path})
    elif "volumes" in self.config:
      self.volumes_list = self.config["volumes"]
    else:
      self.volumes_list = None
      

    if(not self.volumes_list is None):
      gcs_volume = batch_v1.Volume()
      for volume in self.volumes_list:
        gcs_bucket = batch_v1.GCS()
        gcs_bucket.remote_path = volume["bucket_name"]
        gcs_volume.mount_path = volume["gcs_path"]
        gcs_volume.gcs = gcs_bucket
        self.task.volumes.append(gcs_volume)
        if "container" in self.config:
          self.runnable.container.volumes.append(
            volume["gcs_path"]+":"+volume["gcs_path"]+":rw")

  # We can specify what resources are requoested by each task.

    resources = batch_v1.ComputeResource()
    resources.cpu_milli = self.config["cpu_milli"] if "cpu_milli" in self.config else 1000 
    resources.memory_mib = self.config["memory_mib"] if "memory_mib" in self.config else 102400
    resources.boot_disk_mib = self.config["boot_disk_mib"] if "boot_disk_mib" in self.config else 102400
    self.task.compute_resource = resources

    return(self.task)

    
  def _create_allocation_policy(self) -> batch_v1.AllocationPolicy:

      # Policies are used to define on what kind of virtual machines the tasks will run on.
      # In this case, we tell the system to use an instance template that defines all the
      # required parameters.

      self.allocation_policy = batch_v1.AllocationPolicy()
      self.instance = batch_v1.AllocationPolicy.InstancePolicyOrTemplate()
      self.instance_policy = batch_v1.AllocationPolicy.InstancePolicy()
      self.accelerator = batch_v1.AllocationPolicy.Accelerator()
      self.network_policy = batch_v1.AllocationPolicy.NetworkPolicy()
      self.network = batch_v1.AllocationPolicy.NetworkInterface()

      if "template_link" in self.config:
        self.instance.instance_template = self.config["template_link"]
      elif "machine_type" in self.config:
        self.instance_policy.machine_type = self.config["machine_type"]
        self.instance.policy = self.instance_policy

        if "accelerator" in self.config:
          self.accelerator.type_ = self.config["accelerator"]["type"]
          self.accelerator.count = self.config["accelerator"]["count"]
          self.instance_policy.accelerators = [self.accelerator]
        if "install_gpu_drivers" in self.config:
          self.instance.install_gpu_drivers = True
        self.instance.policy = self.instance_policy
          
      else:
        raise(Error("No instance policy defined."))

      location_policy = batch_v1.AllocationPolicy.LocationPolicy()
      self.allocation_policy.instances = [self.instance]

      if "network" in self.config:
        if "no_external_ip_address" in self.config:
          self.network.no_external_ip_address = self.config["no_external_ip_address"]
          self.network.subnetwork = f'regions/{self.config["region"] }/subnetworks/{self.config["subnetwork"]}'
          
        self.network.network = f'projects/{ self.project_id }/global/networks/{self.config["network"]}'
        self.network_policy.network_interfaces = [self.network]
        self.allocation_policy.network = self.network_policy

      if "allowed_locations" in self.config:
        location_policy.allowed_locations = self.config["allowed_locations"]
        self.allocation_policy.location = location_policy

      return(self.allocation_policy)
    
  def _create_taskgroup(self):
      """ 
        Tasks are grouped inside a job using TaskGroups.
        Currently, it's possible to have only one task group.
      """

      self.group = batch_v1.TaskGroup()
      self.group.task_spec = self.task
      self.group.task_count = self.config["task_count"] if "task_count" in self.config else 1
      self.group.parallelism = self.config["parallelism"] if "parallelism" in self.config else 1
      self.group.task_count_per_node = self.config["task_count_per_node"] if "task_count_per_node" in self.config else 1

      return(self.group)

  def _create_job_request(self):
      """Creates job after configuration is completed. 
      """

      self.job = batch_v1.Job()
      self.job.labels = self.config["labels"] if "labels" in self.config else {"env": "hpc", "type": "hpc"}

      self.job.task_groups = [self.group]
      self.job.allocation_policy = self.allocation_policy

      # We use Cloud Logging as it's an out of the box available option
      self.job.logs_policy = batch_v1.LogsPolicy()
      self.job.logs_policy.destination = batch_v1.LogsPolicy.Destination.CLOUD_LOGGING
      return(self.job)
    
  def create_job_request(self) -> batch_v1.Job:

      # Define what will be done as part of the job.

      self._create_runnable()
      self._create_task()
      self.task.runnables = [self.runnable]
      self._create_taskgroup()
      self._create_allocation_policy()
      self._create_job_request()

      create_request = batch_v1.CreateJobRequest()
      create_request.job = self.job
      create_request.job_id = self.job_id
      
      # The job's parent is the region in which the job will run
      create_request.parent = f'projects/{ self.project_id }/locations/{self.config["region"]}'

      return(create_request)
  # [END batch_create_job_with_template]

  def delete_job(self, job_to_delete) -> Operation:
      """
      Triggers the deletion of a Job.

      Args:
        delete_job: name of the job to delete

      Returns:
          An operation object related to the deletion. You can call `.result()`
          on it to wait for its completion.
      """
      client = batch_v1.BatchServiceClient()

      return client.delete_job(name=f"projects/{self.project_id}/locations/{self.config['region']}/jobs/{job_to_delete}")

  def list_jobs(self) -> Iterable[batch_v1.Job]:
      """
      Get a list of all jobs defined in given region.

      Returns:
          An iterable collection of Job object.
      """
      client = batch_v1.BatchServiceClient()

      return client.list_jobs(parent=f"projects/{self.project_id}/locations/{self.config['region']}")



def parse_yaml_file(file_name: str):
  """
  Parse the provided YAML file to get the job configuration

  :param file_name: The path to the YAML configuration file
  :type file_name: str
  """
  with open(file_name) as f:
    data = yaml.load(f, Loader=SafeLoader)
    return(data)


def main(argv):
  """
  This is where the Job and Pubsub objects are created and used.

  :param argv: Standard commandline arguments
  :type argv: _type_
  """

  # Read in config file. 
  config = parse_yaml_file(FLAGS.config_file)

  # Create unique job ID, required for Batch and k8s Jobs
  job_id = config["job_prefix"] + uuid.uuid4().hex[:8]

  # "previous_job_id": Used for Pubsub. TOPIC is created from JobID, a restart can provide an existing Job
  # Id for existing pubsub queue
  # TODO: support more env_vars

  env_vars = {}
  if FLAGS.previous_job_id:
    env_vars = dict(env_vars, TOPIC_ID = FLAGS.previous_job_id)
  else:
    env_vars = dict(env_vars, TOPIC_ID = job_id)

  # Create jobs object
  jobs = CloudBatchJobs(job_id, config, env_vars)


  # If pubsub queue is required 
  if FLAGS.pubsub and not FLAGS.debug:

    pubsub = PubSub(job_id, config, previous_job_id=FLAGS.previous_job_id)
    pubsub.create_topic()
    pubsub.create_subscription()
    pubsub.publish_fifo_ids()

  # Delete job. JobID must be passed.
  if(FLAGS.delete_job):
    print("Deleting job", file=sys.stderr)
    deleted_job = jobs.delete_job(FLAGS.delete_job)
    print(deleted_job.result, file=sys.stderr)
    exit()

  # Prints list of jobs, in queue, running or complteted.
  if(FLAGS.list_jobs):
    print("Listing jobs", file=sys.stderr)
    job_list = jobs.list_jobs()

    for job in job_list:
      print(job.name,"\t",job.status.state, file=sys.stderr)
    exit()

  # Prints 
  if FLAGS.debug:
    print(config, file=sys.stderr)
    
  if FLAGS.create_job:
    # Create the job
    print(jobs.client.create_job(jobs.create_job_request()), file=sys.stderr)
  else:
    print(jobs.create_job_request().job, file=sys.stderr)

if __name__ == "__main__":
    """ This is executed when run from the command line """
    app.run(main)