Fixes for resource deployment and added covalent version specification (

#30) * minor fix for azure batch plugin * Added covalent_package_version to Dockerfile * updated changelog * removed cloudpickle requirement * installing GCC in docker * resource spinning up fixed * storing client secret in as well * updated changelog * Update CHANGELOG.md Co-authored-by: Will Cunningham <[email protected]> --------- Co-authored-by: Manjunath PV <[email protected]> Co-authored-by: Will Cunningham <[email protected]>
AgnostiqHQ · Jan 17, 2024 · 62e50d1 · 62e50d1
1 parent 9737d13
commit 62e50d1
Show file tree

Hide file tree

Showing 10 changed files with 102 additions and 21 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### Added
+
+- Added `covalent_package_version` so that we can specify the version of the covalent to use when building the docker image
+- Added `covalent_package_version`, `create_batch_acount`, `batch_account_name`, and `batch_resource_group` to `ExecutorInfraDefaults` to accept arguments from the `covalent deploy` command
+- Added `_EXECUTOR_PLUGIN_DEFAULTS` variable to populate Azure Batch plugin defaults in the `covalent` config file
+- Added `.gitignore` file in the `infra` directory to ignore more unnecessary files from being committed
+
+### Fixed
+
+- Added `subscription_id`, `owners`, `tenant_id`, to ExecutorInfraDefaults to accept arguments from the `covalent deploy` command
+- Fixed `pool_id` value in config file
+- Fixed `gcc` installation in the docker image
+
+### Changed
+
+- Modified the `owners` variable to have a default value
+- `base_image_uri` and `client_secret` are now also saved in the `covalent` config file so that the user doesn't have to manually go and look for their values in the tfstate file
+
 ## [0.16.0] - 2023-11-21
 
 ### Changed

diff --git a/covalent_azurebatch_plugin/assets/infra/.gitignore b/covalent_azurebatch_plugin/assets/infra/.gitignore
@@ -0,0 +1,4 @@
+*.conf
+.docker
+*.sh
+*.log
diff --git a/covalent_azurebatch_plugin/assets/infra/Dockerfile b/covalent_azurebatch_plugin/assets/infra/Dockerfile
@@ -17,15 +17,23 @@
 ARG COVALENT_BASE_IMAGE
 FROM ${COVALENT_BASE_IMAGE}
 
+ARG COVALENT_PACKAGE_VERSION
+
 USER root
 
-RUN apt-get update && apt-get install -y \
+RUN apt-get update && apt-get install gcc -y \
   && rm -rf /var/lib/apt/lists/* \
   && pip install --no-cache-dir --upgrade \
   azure-batch==12.0.0 \
   azure-identity==1.11.0 \
-  azure-storage-blob==12.14.1 \
-  "covalent>=0.202.0,<1"
+  azure-storage-blob==12.14.1
+
+
+RUN if [ -z "$COVALENT_PACKAGE_VERSION" ]; then \
+		pip install covalent; else \
+		pip install "covalent==$COVALENT_PACKAGE_VERSION"; \
+	fi
+
 
 WORKDIR /covalent
 

diff --git a/covalent_azurebatch_plugin/assets/infra/azurebatch.conf.tftpl b/covalent_azurebatch_plugin/assets/infra/azurebatch.conf.tftpl
@@ -1,9 +1,11 @@
 [azurebatch]
 tenant_id = ${tenant_id}
 client_id = ${client_id}
+client_secret = ${client_secret}
 batch_account_url = ${batch_account_url}
 batch_account_domain = ${batch_account_domain}
 storage_account_name = ${storage_account_name}
 storage_account_domain = ${storage_account_domain}
-pool_id = ${storage_account_domain}
+pool_id = ${pool_id}
 retries = ${retries}
+base_image_uri = ${base_image_uri}
diff --git a/covalent_azurebatch_plugin/assets/infra/iam.tf b/covalent_azurebatch_plugin/assets/infra/iam.tf
@@ -20,8 +20,8 @@
 
 resource "azurerm_user_assigned_identity" "batch" {
   name                = "${var.prefix}covalentbatch"
-  resource_group_name = azurerm_resource_group.batch.name
-  location            = azurerm_resource_group.batch.location
+  resource_group_name = var.create_batch_account ? azurerm_resource_group.batch[0].name : data.azurerm_resource_group.batch[0].name
+  location            = var.region
 }
 
 resource "azurerm_role_assignment" "batch_to_acr" {

diff --git a/covalent_azurebatch_plugin/assets/infra/main.tf b/covalent_azurebatch_plugin/assets/infra/main.tf
@@ -27,23 +27,37 @@ provider "azurerm" {
 
 resource "azurerm_resource_group" "batch" {
   name     = "${var.prefix}-covalent-batch"
+  count    = var.create_batch_account ? 1 : 0
   location = var.region
 }
 
+data "azurerm_resource_group" "batch" {
+  name =  var.batch_resource_group
+  count = var.create_batch_account ? 0 : 1
+}
+
 resource "azurerm_batch_account" "covalent" {
   name                = "${var.prefix}covalentbatch"
-  resource_group_name = azurerm_resource_group.batch.name
-  location            = azurerm_resource_group.batch.location
+  count               = var.create_batch_account ? 1 : 0
+  resource_group_name = azurerm_resource_group.batch[0].name
+  location            = azurerm_resource_group.batch[0].location
 
   storage_account_id                  = azurerm_storage_account.batch.id
   storage_account_authentication_mode = "StorageKeys"
 }
 
+data "azurerm_batch_account" "covalent" {
+  name                = var.batch_account_name
+  count               = var.create_batch_account ? 0 : 1
+  resource_group_name = data.azurerm_resource_group.batch[0].name
+}
+
+
 resource "azurerm_batch_pool" "covalent" {
-  name                = "default"
-  resource_group_name = azurerm_resource_group.batch.name
+  name                = "${var.prefix}-default"
+  resource_group_name = var.create_batch_account ? azurerm_resource_group.batch[0].name : data.azurerm_resource_group.batch[0].name
 
-  account_name = azurerm_batch_account.covalent.name
+  account_name = var.create_batch_account ? azurerm_batch_account.covalent[0].name : data.azurerm_batch_account.covalent[0].name
   display_name = "Covalent Azure Plugin Default Pool"
 
   vm_size           = var.vm_name
@@ -91,13 +105,15 @@ data "template_file" "executor_config" {
   vars = {
     subscription_id        = var.subscription_id
     tenant_id              = var.tenant_id
-    client_id              = "${azuread_application.batch.application_id}"
-    batch_account_url      = "https://${azurerm_batch_account.covalent.account_endpoint}"
+    client_id              = "${azuread_application.batch.client_id}"
+    client_secret          = "${azuread_service_principal_password.covalent_plugin.value}"
+    batch_account_url      = var.create_batch_account ? "https://${azurerm_batch_account.covalent[0].account_endpoint}" : "https://${data.azurerm_batch_account.covalent[0].account_endpoint}"
     batch_account_domain   = "batch.core.windows.net"
     storage_account_name   = "${azurerm_storage_account.batch.name}"
     storage_account_domain = "blob.core.windows.net"
     pool_id                = "${azurerm_batch_pool.covalent.name}"
     retries                = 3
+    base_image_uri         = "${azurerm_container_registry.batch.login_server}/covalent-executor-base:latest"
   }
 }
 

diff --git a/covalent_azurebatch_plugin/assets/infra/outputs.tf b/covalent_azurebatch_plugin/assets/infra/outputs.tf
@@ -23,7 +23,7 @@ output "user_identity_resource_id" {
 }
 
 output "plugin_client_username" {
-  value = azuread_application.batch.application_id
+  value = azuread_application.batch.client_id
 }
 
 output "plugin_client_secret" {
@@ -37,7 +37,7 @@ output "covalent_azurebatch_object" {
         tenant_id="${var.tenant_id}",
         client_id="${azuread_application.batch.application_id}",
         client_secret=plugin_client_secret,
-        batch_account_url="https://${azurerm_batch_account.covalent.account_endpoint}",
+        batch_account_url="https://${var.create_batch_account ? azurerm_batch_account.covalent[0].account_endpoint : data.azurerm_batch_account.covalent[0].account_endpoint}",
         storage_account_name="${azurerm_storage_account.batch.name}",
         pool_id="${azurerm_batch_pool.covalent.name}",
     )

diff --git a/covalent_azurebatch_plugin/assets/infra/storage.tf b/covalent_azurebatch_plugin/assets/infra/storage.tf
@@ -16,16 +16,16 @@
 
 resource "azurerm_container_registry" "batch" {
   name                = "${var.prefix}covalentbatch"
-  resource_group_name = azurerm_resource_group.batch.name
-  location            = azurerm_resource_group.batch.location
+  resource_group_name = var.create_batch_account ? azurerm_resource_group.batch[0].name : data.azurerm_resource_group.batch[0].name
+  location            = var.region
 
   sku = "Basic"
 
   provisioner "local-exec" {
     interpreter = ["/bin/bash", "-c"]
     command     = <<EOL
       set -eu -o pipefail
-      docker build --build-arg COVALENT_BASE_IMAGE=python:3.8-slim-bullseye --tag ${self.login_server}/covalent-executor-base:latest .
+      docker build --no-cache --build-arg COVALENT_BASE_IMAGE=python:3.8-slim-bullseye --build-arg COVALENT_PACKAGE_VERSION=${var.covalent_package_version} --tag ${self.login_server}/covalent-executor-base:latest .
       az acr login --name ${self.name}
       docker push -a ${self.login_server}/covalent-executor-base
     EOL
@@ -34,8 +34,8 @@ resource "azurerm_container_registry" "batch" {
 
 resource "azurerm_storage_account" "batch" {
   name                = "${var.prefix}covalentbatch"
-  resource_group_name = azurerm_resource_group.batch.name
-  location            = azurerm_resource_group.batch.location
+  resource_group_name = var.create_batch_account ? azurerm_resource_group.batch[0].name : data.azurerm_resource_group.batch[0].name
+  location            = var.region
 
   account_tier             = "Standard"
   account_replication_type = "LRS"

diff --git a/covalent_azurebatch_plugin/assets/infra/variables.tf b/covalent_azurebatch_plugin/assets/infra/variables.tf
@@ -40,9 +40,34 @@ variable "subscription_id" {
 variable "owners" {
   description = "List of owner IDs for the service principal credentials"
   type        = list(string)
+  default     = []
 }
 
 variable "vm_name" {
   description = "Name of the VM used in the Batch pool"
   default     = "Standard_A1_v2"
 }
+
+variable "covalent_package_version" {
+  type        = string
+  description = "Covalent version to be installed in the container, if not specified the latest stable version will be installed"
+  default     = ""
+}
+
+variable "create_batch_account" {
+  type        = bool
+  description = "Whether to create a new Batch account or use an existing one"
+  default     = true
+}
+
+variable "batch_account_name" {
+  type        = string
+  description = "Name of the Batch account to be used"
+  default     = ""
+}
+
+variable "batch_resource_group" {
+  type        = string
+  description = "Name of the resource group containing the Batch account to be used"
+  default     = ""
+}
diff --git a/covalent_azurebatch_plugin/azurebatch.py b/covalent_azurebatch_plugin/azurebatch.py
@@ -66,7 +66,9 @@ class ExecutorInfraDefaults(BaseModel):
     """
 
     prefix: Optional[str] = "covalent-batch"
-    tenant_id: str = ""
+    subscription_id: str
+    owners: List[str] = []
+    tenant_id: str
     client_id: str = ""
     client_secret: str = ""
     batch_account_url: str = ""
@@ -81,10 +83,16 @@ class ExecutorInfraDefaults(BaseModel):
     time_limit: int = 300
     cache_dir: str = "/tmp/covalent"
     poll_freq: int = 30
+    covalent_package_version: str = ""
+    create_batch_account: bool = True
+    batch_account_name: str = ""
+    batch_resource_group: str = ""
 
 
 EXECUTOR_PLUGIN_NAME = "AzureBatchExecutor"
 
+_EXECUTOR_PLUGIN_DEFAULTS = ExecutorPluginDefaults().dict()
+
 FUNC_FILENAME = "func-{dispatch_id}-{node_id}.pkl"
 RESULT_FILENAME = "result-{dispatch_id}-{node_id}.pkl"
 STORAGE_CONTAINER_NAME = "covalent-assets"