def _create_allocation_policy()

in python-batch/batch.py [0:0]

• 34 lines of code
• 8 McCabe index (conditional complexity)

  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)