From f8c811be16ec63b92436a1a5b500700f04d13072 Mon Sep 17 00:00:00 2001
From: Sankalp Sanand <sankalp@agnostiq.ai>
Date: Mon, 8 Jan 2024 14:47:19 -0500
Subject: [PATCH] Minor fixes for dispatching issues (#35)

* made some minor fixes to resolve dispatching issues

* updated changelog

* removed cloudpickle upperbound

* updated changelog

* dict to model_dump

* dict to model_dump reverted for now
---
 CHANGELOG.md                                       |  8 ++++++++
 covalent_gcpbatch_plugin/assets/docker/Dockerfile  | 14 +++++++++++---
 covalent_gcpbatch_plugin/assets/infra/variables.tf |  4 ++--
 covalent_gcpbatch_plugin/gcpbatch.py               |  1 +
 4 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b679290..b7fa5ff 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### Changed
+
+- No longer using `PRE_RELEASE` variable, instead switched to directly specifying the covalent version number to install. This was required to support its passing in the new `covalent deploy` command.
+
+### Added
+
+- `covalent_package_version` can now be specified directly from the `covalent deploy` command.
+
 ## [0.14.0] - 2023-12-05
 
 ### Authors
diff --git a/covalent_gcpbatch_plugin/assets/docker/Dockerfile b/covalent_gcpbatch_plugin/assets/docker/Dockerfile
index 07efd7e..77bf0cf 100644
--- a/covalent_gcpbatch_plugin/assets/docker/Dockerfile
+++ b/covalent_gcpbatch_plugin/assets/docker/Dockerfile
@@ -21,14 +21,22 @@ FROM ${COVALENT_BASE_IMAGE}
 ARG COVALENT_TASK_ROOT=/usr/src
 
 ARG COVALENT_PACKAGE_VERSION
+
+# Not used anymore for now
 ARG PRE_RELEASE
 
 COPY requirements-image.txt requirements.txt
 RUN apt-get update && pip install -r requirements.txt
 
-RUN if [ -z "$PRE_RELEASE" ]; then \
-		pip install "$COVALENT_PACKAGE_VERSION"; else \
-		pip install --pre "$COVALENT_PACKAGE_VERSION"; \
+# Commenting this out for now as we switch to specifying the package version directly
+# RUN if [ -z "$PRE_RELEASE" ]; then \
+# 		pip install "$COVALENT_PACKAGE_VERSION"; else \
+# 		pip install --pre "$COVALENT_PACKAGE_VERSION"; \
+# 	fi
+
+RUN if [ -z "$COVALENT_PACKAGE_VERSION" ]; then \
+		pip install covalent; else \
+		pip install "covalent==$COVALENT_PACKAGE_VERSION"; \
 	fi
 
 
diff --git a/covalent_gcpbatch_plugin/assets/infra/variables.tf b/covalent_gcpbatch_plugin/assets/infra/variables.tf
index 5c56f38..4a3e7c4 100644
--- a/covalent_gcpbatch_plugin/assets/infra/variables.tf
+++ b/covalent_gcpbatch_plugin/assets/infra/variables.tf
@@ -28,8 +28,8 @@ variable "prerelease" {
 
 variable "covalent_package_version" {
   type        = string
-  description = "Covalent version to be installed in the container"
-  default     = "covalent"
+  description = "Covalent version to be installed in the container, if not specified the latest stable version will be installed"
+  default     = ""
 }
 
 variable "prefix" {
diff --git a/covalent_gcpbatch_plugin/gcpbatch.py b/covalent_gcpbatch_plugin/gcpbatch.py
index fdac77d..1c78796 100644
--- a/covalent_gcpbatch_plugin/gcpbatch.py
+++ b/covalent_gcpbatch_plugin/gcpbatch.py
@@ -64,6 +64,7 @@ class ExecutorInfraDefaults(BaseModel):
     retries: Optional[int] = 3
     cache_dir: Optional[str] = "/tmp/covalent"
     key_path: str
+    covalent_package_version: str = ""
 
 
 _EXECUTOR_PLUGIN_DEFAULTS = ExecutorPluginDefaults().dict()