fix: adapt to changes in Snakemake 8

johanneskoester · johanneskoester · commit ca853cb80529 · 2023-10-31T09:53:48.000+01:00
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -13,15 +13,19 @@ jobs:
       - name: Check out the code
         uses: actions/checkout@v3
 
+      - uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+
       - name: Install poetry
-        run: pipx install poetry
+        run: pip install poetry
 
       - name: Determine dependencies
         run: poetry lock
 
       - uses: actions/setup-python@v4
         with:
-          python-version: "3.9"
+          python-version: "3.11"
           cache: poetry
 
       - name: Install Dependencies using Poetry
@@ -36,15 +40,19 @@ jobs:
       - name: Check out the code
         uses: actions/checkout@v3
 
+      - uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+
       - name: Install poetry
-        run: pipx install poetry
+        run: pip install poetry
 
       - name: Determine dependencies
         run: poetry lock
 
       - uses: actions/setup-python@v4
         with:
-          python-version: "3.9"
+          python-version: "3.11"
           cache: poetry
 
       - name: Install Dependencies using Poetry
@@ -58,26 +66,26 @@ jobs:
     steps:
       - uses: actions/checkout@v3
 
+      - uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+
       - name: Install poetry
-        run: pipx install poetry
+        run: pip install poetry
 
       - name: Determine dependencies
         run: poetry lock
 
       - uses: actions/setup-python@v4
         with:
-          python-version: "3.9"
+          python-version: "3.11"
           cache: poetry
 
-      - name: Install dependencies
-        run: |
-          poetry install
+      - name: Install Dependencies using Poetry
+        run: poetry install
 
       - name: Run pytest
         env:
-          AZ_BLOB_PREFIX: "${{ secrets.AZ_BLOB_PREFIX }}"
-          AZ_BLOB_ACCOUNT_URL: "${{ secrets.AZ_STORAGE_ACCOUNT_URL }}"
-          AZ_BLOB_CREDENTIAL: "${{ secrets.AZ_STORAGE_KEY }}"
           AZ_BATCH_ACCOUNT_URL: "${{ secrets.AZ_BATCH_ACCOUNT_URL }}"
           AZ_BATCH_ACCOUNT_KEY: "${{ secrets.AZ_BATCH_KEY }}"
         run: poetry run coverage run -m pytest tests/tests.py
diff --git a/pyproject.toml b/pyproject.toml
@@ -6,9 +6,9 @@ authors = ["Johannes Koester <johannes.koester@uni-due.de>"]
 readme = "README.md"
 
 [tool.poetry.dependencies]
-python = "^3.9"
-snakemake-interface-common = "^1.12.0"
-snakemake-interface-executor-plugins = "^6.0.0"
+python = "^3.11"
+snakemake-interface-common = "^1.13.0"
+snakemake-interface-executor-plugins = "^7.0.0"
 azure-storage-blob = "^12.17.0"
 azure-batch = "^14.0.0"
 azure-mgmt-batch = "^17.1.0"
@@ -21,7 +21,7 @@ flake8 = "^6.1.0"
 coverage = "^7.3.1"
 pytest = "^7.4.2"
 snakemake = {git = "https://github.com/snakemake/snakemake.git"}
-snakemake-storage-plugin-s3 = "^0.2.0"
+snakemake-storage-plugin-s3 = "^0.2.1"
 
 [build-system]
 requires = ["poetry-core"]
diff --git a/snakemake_executor_plugin_azure_batch/__init__.py b/snakemake_executor_plugin_azure_batch/__init__.py
@@ -17,13 +17,11 @@
 import uuid
 from pprint import pformat
 from urllib.parse import urlparse
-from typing import List, Generator, Optional
+from typing import AsyncGenerator, List, Optional
 
 from snakemake_interface_executor_plugins.executors.base import SubmittedJobInfo
 from snakemake_interface_executor_plugins.executors.remote import RemoteExecutor
 from snakemake_interface_executor_plugins import ExecutorSettingsBase, CommonSettings
-from snakemake_interface_executor_plugins.workflow import WorkflowExecutorInterface
-from snakemake_interface_executor_plugins.logging import LoggerExecutorInterface
 from snakemake_interface_executor_plugins.jobs import (
     JobExecutorInterface,
 )
@@ -42,7 +40,6 @@
 
 from snakemake.remote.AzBlob import AzureStorageHelper
 import msrest.authentication as msa
-from snakemake_executor_plugin_azure_batch.common import bytesto
 
 
 # Optional:
@@ -91,29 +88,17 @@ class ExecutorSettings(ExecutorSettingsBase):
     # filesystem (True) or not (False).
     # This is e.g. the case for cloud execution.
     implies_no_shared_fs=True,
+    pass_default_storage_provider_args=True,
+    pass_default_resources_args=True,
+    pass_envvar_declarations_to_cmd=False,
+    auto_deploy_default_storage_provider=True,
 )
 
 
 # Required:
 # Implementation of your executor
 class Executor(RemoteExecutor):
-    def __init__(
-        self,
-        workflow: WorkflowExecutorInterface,
-        logger: LoggerExecutorInterface,
-    ):
-        super().__init__(
-            workflow,
-            logger,
-            # configure behavior of RemoteExecutor below
-            # whether arguments for setting the remote provider shall  be passed to jobs
-            pass_default_storage_provider_args=True,
-            # whether arguments for setting default resources shall be passed to jobs
-            pass_default_resources_args=True,
-            # whether environment variables shall be passed to jobs
-            pass_envvar_declarations_to_cmd=True,
-        )
-
+    def __post_init__(self):
         AZURE_BATCH_RESOURCE_ENDPOINT = "https://batch.core.windows.net/"
 
         # Here we validate that az blob credential is SAS
@@ -124,14 +109,12 @@ def __init__(
         # TODO this does not work if the remote is used without default_remote_prefix
         # get container from remote prefix
         self.prefix_container = str.split(
-            workflow.storage_settings.default_remote_prefix, "/"
+            self.workflow.storage_settings.default_remote_prefix, "/"
         )[0]
 
         # setup batch configuration sets self.az_batch_config
         self.batch_config = AzBatchConfig(self.workflow.executor_settings.account_url)
-        logger.debug(f"AzBatchConfig: {self.mask_batch_config_as_string()}")
-
-        self.workflow = workflow
+        self.logger.debug(f"AzBatchConfig: {self.mask_batch_config_as_string()}")
 
         # handle case on OSX with /var/ symlinked to /private/var/ causing
         # issues with workdir not matching other workflow file dirs
@@ -156,28 +139,19 @@ def __init__(
         # enable autoscale flag
         self.az_batch_enable_autoscale = self.workflow.executor_settings.autoscale
 
-        # Package workflow sources files and upload to storage
-        self._build_packages = set()
-        targz = self._generate_build_source_package()
-
-        # removed after job failure/success
-        self.resource_file = self._upload_build_source_package(
-            targz, resource_prefix=self.batch_config.resource_file_prefix
-        )
-
         # authenticate batch client from SharedKeyCredentials
         if (
             self.batch_config.batch_account_key is not None
             and self.batch_config.managed_identity_client_id is None
         ):
-            logger.debug("Using batch account key for authentication...")
+            self.logger.debug("Using batch account key for authentication...")
             creds = SharedKeyCredentials(
                 self.batch_config.batch_account_name,
                 self.batch_config.batch_account_key,
             )
         # else authenticate with managed indentity client id
         elif self.batch_config.managed_identity_client_id is not None:
-            logger.debug("Using managed identity batch authentication...")
+            self.logger.debug("Using managed identity batch authentication...")
             creds = DefaultAzureCredential(
                 managed_identity_client_id=self.batch_config.managed_identity_client_id
             )
@@ -192,7 +166,7 @@ def __init__(
         if self.batch_config.managed_identity_resource_id is not None:
             self.batch_mgmt_client = BatchManagementClient(
                 credential=DefaultAzureCredential(
-                    managed_identity_client_id=self.batch_config.managed_identity_client_id # noqa
+                    managed_identity_client_id=self.batch_config.managed_identity_client_id  # noqa
                 ),
                 subscription_id=self.batch_config.subscription_id,
             )
@@ -210,12 +184,6 @@ def shutdown(self):
         self.logger.debug("Deleting AzBatch pool")
         self.batch_client.pool.delete(self.pool_id)
 
-        self.logger.debug("Deleting workflow sources from blob")
-
-        self.azblob_helper.delete_from_container(
-            self.prefix_container, self.resource_file.file_path
-        )
-
         super().shutdown()
 
     def run_job(self, job: JobExecutorInterface):
@@ -238,10 +206,7 @@ def run_job(self, job: JobExecutorInterface):
                 continue
 
         exec_job = self.format_job_exec(job)
-        exec_job = (
-            f"/bin/bash -c 'tar xzf {self.resource_file.file_path} && "
-            f"{shlex.quote(exec_job)}'"
-        )
+        exec_job = f"/bin/bash -c '{shlex.quote(exec_job)}'"
 
         # A string that uniquely identifies the Task within the Job.
         task_uuid = str(uuid.uuid1())
@@ -267,7 +232,6 @@ def run_job(self, job: JobExecutorInterface):
             id=task_id,
             command_line=exec_job,
             container_settings=task_container_settings,
-            resource_files=[self.resource_file],  # Snakefile, envs, yml files etc.
             user_identity=batchmodels.UserIdentity(auto_user=user),
             environment_settings=envsettings,
         )
@@ -284,7 +248,7 @@ def run_job(self, job: JobExecutorInterface):
 
     async def check_active_jobs(
         self, active_jobs: List[SubmittedJobInfo]
-    ) -> Generator[SubmittedJobInfo, None, None]:
+    ) -> AsyncGenerator[SubmittedJobInfo, None, None]:
         # Check the status of active jobs.
 
         # You have to iterate over the given list active_jobs.
@@ -321,14 +285,6 @@ async def check_active_jobs(
                     )
                 )
 
-                def print_output():
-                    self.logger.debug(
-                        "task {}: stderr='{}'\n".format(batch_job.task_id, stderr)
-                    )
-                    self.logger.debug(
-                        "task {}: stdout='{}'\n".format(batch_job.task_id, stdout)
-                    )
-
                 if (
                     task.execution_info.result
                     == batchmodels.TaskExecutionResult.failure
@@ -645,94 +601,6 @@ def validate_az_blob_credential_is_sas():
                     "AZ_BLOB_CREDENTIAL is not a valid storage account SAS token."
                 )
 
-    # from google_lifesciences.py
-    def _set_workflow_sources(self):
-        """We only add files from the working directory that are config related
-        (e.g., the Snakefile or a config.yml equivalent), or checked into git.
-        """
-        self.workflow_sources = []
-
-        for wfs in self.dag.get_sources():
-            if os.path.isdir(wfs):
-                for dirpath, dirnames, filenames in os.walk(wfs):
-                    self.workflow_sources.extend(
-                        [
-                            self._check_source_size(os.path.join(dirpath, f))
-                            for f in filenames
-                        ]
-                    )
-            else:
-                self.workflow_sources.append(
-                    self._check_source_size(os.path.abspath(wfs))
-                )
-
-    # from google_lifesciences.py
-    def _generate_build_source_package(self):
-        """in order for the instance to access the working directory in storage,
-        we need to upload it. This file is cleaned up at the end of the run.
-        We do this, and then obtain from the instance and extract.
-        """
-        # Workflow sources for cloud executor must all be under same workdir root
-        for filename in self.workflow_sources:
-            if self.workdir not in filename:
-                raise WorkflowError(
-                    "All source files must be present in the working directory, "
-                    "{workdir} to be uploaded to a build package that respects "
-                    "relative paths, but {filename} was found outside of this "
-                    "directory. Please set your working directory accordingly, "
-                    "and the path of your Snakefile to be relative to it.".format(
-                        workdir=self.workdir, filename=filename
-                    )
-                )
-
-        # We will generate a tar.gz package, renamed by hash
-        tmpname = next(tempfile._get_candidate_names())
-        targz = os.path.join(tempfile.gettempdir(), f"snakemake-{tmpname}.tar.gz")
-        tar = tarfile.open(targz, "w:gz")
-
-        # Add all workflow_sources files
-        for filename in self.workflow_sources:
-            arcname = filename.replace(self.workdir + os.path.sep, "")
-            tar.add(filename, arcname=arcname)
-        self.logger.debug(
-            f"Created {targz} with the following contents: {self.workflow_sources}"
-        )
-        tar.close()
-
-        # Rename based on hash, in case user wants to save cache
-        hasher = hashlib.sha256()
-        hasher.update(open(targz, "rb").read())
-        sha256 = hasher.hexdigest()
-
-        hash_tar = os.path.join(
-            self.workflow.persistence.aux_path, f"workdir-{sha256}.tar.gz"
-        )
-
-        # Only copy if we don't have it yet, clean up if we do
-        if not os.path.exists(hash_tar):
-            shutil.move(targz, hash_tar)
-        else:
-            os.remove(targz)
-
-        # We will clean these all up at shutdown
-        self._build_packages.add(hash_tar)
-
-        return hash_tar
-
-    def _upload_build_source_package(self, targz, resource_prefix=""):
-        """given a .tar.gz created for a workflow, upload it to the blob
-        storage account, only if the blob doesn't already exist.
-        """
-        blob_name = os.path.join(resource_prefix, os.path.basename(targz))
-
-        # upload blob to storage using storage helper
-        bc = self.azblob_helper.upload_to_azure_storage(
-            self.prefix_container, targz, blob_name=blob_name
-        )
-
-        # return resource file
-        return batchmodels.ResourceFile(http_url=bc.url, file_path=blob_name)
-
     # from https://github.com/Azure-Samples/batch-python-quickstart/blob/master/src/python_quickstart_client.py # noqa
     @staticmethod
     def _read_stream_as_string(stream, encoding):
@@ -765,19 +633,6 @@ def _get_task_output(self, job_id, task_id, stdout_or_stderr, encoding=None):
 
         return content
 
-    def _check_source_size(self, filename, warning_size_gb=0.2):
-        """A helper function to check the filesize, and return the file
-        to the calling function Additionally, given that we encourage these
-        packages to be small, we set a warning at 200MB (0.2GB).
-        """
-        gb = bytesto(os.stat(filename).st_size, "g")
-        if gb > warning_size_gb:
-            self.logger.warning(
-                f"File {filename} (size {gb} GB) is greater than the {warning_size_gb} "
-                "GB suggested size. Consider uploading larger files to storage first."
-            )
-        return filename
-
 
 class AzBatchConfig:
     def __init__(self, batch_account_url: str):
diff --git a/tests/tests.py b/tests/tests.py
