Integrate dsub

config/config.yaml

@@ -3,8 +3,10 @@
 # The length of the hash used to identify a parameter combination
 hash_length: 7
 
-# Specify the container framework. Current supported versions include 'docker' and
-# 'singularity'. If container_framework is not specified, SPRAS will default to docker.
+# Specify the container framework used by each PRM wrapper. Valid options include:
+# - docker (default if not specified)
+# - singularity -- Also known as apptainer, useful in HPC/HTC environments where docker isn't allowed
+# - dsub -- experimental with limited support, used for running on Google Cloud
 container_framework: docker
 
 # Only used if container_framework is set to singularity, this will unpack the singularity containers

environment.yml

@@ -17,6 +17,12 @@ dependencies:
   - scikit-learn=1.2
   - seaborn=0.12
   - spython=0.2
+  # for dsub
+  - python-dateutil<=2.9.0
+  - pytz<=2024.1
+  - pyyaml<=6.0.1
+  - tenacity<=8.2.3
+  - tabulate<=0.9.0
   # Only required for GraphSpace
   - commonmark=0.9
   - docutils=0.19
@@ -27,3 +33,4 @@ dependencies:
   - pip:
      - graphspace_python==1.3.1
      - sphinx-rtd-theme==2.0.0
+     - dsub==0.4.13

spras/config.py

@@ -62,7 +62,7 @@ def __init__(self, raw_config):
         # __init__ makes clear exactly what is being configured.
         # Directory used for storing output
         self.out_dir = None
-        # Container framework used by PRMs. Valid options are "docker" and "singularity"
+        # Container framework used by PRMs. Valid options are "docker", "dsub", and "singularity"
         self.container_framework = None
         # The container prefix (host and organization) to use for images. Default is "docker.io/reedcompbio"
         self.container_prefix = DEFAULT_CONTAINER_PREFIX
@@ -114,9 +114,11 @@ def process_config(self, raw_config):
         # However, if we get a bad value, we raise an exception.
         if "container_framework" in raw_config:
             container_framework = raw_config["container_framework"].lower()
-            if container_framework not in ("docker", "singularity"):
+            if container_framework not in ("docker", "singularity", "dsub"):
                 msg = "SPRAS was configured to run with an unknown container framework: '" + raw_config["container_framework"] + "'. Accepted values are 'docker' or 'singularity'."
                 raise ValueError(msg)
+            if container_framework == "dsub":
+                print("Warning: 'dsub' framework integration is experimental and may not be fully supported.")
             self.container_framework = container_framework
         else:
             self.container_framework = "docker"

spras/containers.py

@@ -1,6 +1,7 @@
 import os
 import platform
 import re
+import subprocess
 from pathlib import Path, PurePath, PurePosixPath
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -41,6 +42,79 @@ def convert_docker_path(src_path: PurePath, dest_path: PurePath, file_path: Unio
     rel_path = file_path.relative_to(src_path)
     return PurePosixPath(dest_path, rel_path)
 
+def download_gcs(gcs_path: str, local_path: str, is_dir: bool):
+    # check that output path exists
+    if not os.path.exists(Path(local_path).parent):
+        os.makedirs(Path(local_path).parent)
+
+    # build command
+    cmd = 'gcloud storage'
+    # rsync with checksums to make file transfer faster for larger files
+    cmd = cmd + ' rsync --checksums-only'
+    # check if directory
+    if is_dir:
+        cmd = cmd + ' -r'
+    cmd = cmd + ' ' + gcs_path + ' ' + local_path
+
+    print(cmd)
+    # run command
+    subprocess.run(cmd, shell=True)
+
+    if is_dir and Path(Path(local_path)/'gcs_temp.txt').exists():
+        os.remove(Path(Path(local_path)/'gcs_temp.txt'))
+
+def upload_gcs(local_path: str, gcs_path: str, is_dir: bool):
+    # check if path exists in cloud storage
+    exists = len(subprocess.run(f'gcloud storage ls {gcs_path}', shell=True, capture_output=True, text=True).stdout)
+    # if path exists rsyc
+    if exists > 0:
+        cmd = 'gcloud storage rsync --checksums-only'
+    # if directory is empty
+    elif exists == 0 and len(os.listdir(local_path)) == 0:
+        # create a temporary file because GCS will not recognize empty directories
+        Path(Path(local_path)/'gcs_temp.txt').touch()
+        # copy path to cloud storage
+        cmd = 'gcloud storage cp -c'
+    # else copy path to cloud storage
+    else:
+        cmd = 'gcloud storage cp -c'
+    # check if directory
+    if is_dir:
+        cmd = cmd + ' -r'
+    cmd = cmd + ' ' + str(Path(local_path).resolve()) + ' ' + gcs_path
+
+    print(cmd)
+    # run command
+    subprocess.run(cmd, shell=True)
+
+def prepare_dsub_cmd(flags: dict):
+    # set constant flags
+    dsub_command = 'dsub'
+    flags['provider'] = 'google-cls-v2'
+    flags['regions'] = 'us-central1'
+    flags['user-project'] = os.getenv('GOOGLE_PROJECT')
+    flags['project'] = os.getenv('GOOGLE_PROJECT')
+    flags['network'] = 'network'
+    flags['subnetwork'] = 'subnetwork'
+    flags['service-account'] = subprocess.run(['gcloud', 'config' ,'get-value' ,'account'], capture_output=True, text=True).stdout.replace('\n', '')
+
+    # order flags according to flag_list
+    flag_list = ["provider", "regions", "zones", "location", "user-project", "project", "network", "subnetwork", "service-account", "image", "env", "logging", "input", "input-recursive", "mount", "output", "output-recursive", "command", "script"]
+    ordered_flags = {f:flags[f] for f in flag_list if f in flags.keys()}
+
+    # iteratively add flags to the command
+    for flag in ordered_flags.keys():
+        if isinstance(ordered_flags.get(flag), list):
+            for f in ordered_flags.get(flag):
+                dsub_command = dsub_command + " --" + flag + " " + f
+        else:
+            dsub_command = dsub_command + " --" + flag + " " + ordered_flags.get(flag)
+
+    # Wait for dsub job to complegte
+    dsub_command = dsub_command + " --wait"
+    print(f"Command: {dsub_command}")
+    return dsub_command
+
 
 # TODO consider a better default environment variable
 # TODO environment currently a single string (e.g. 'TMPDIR=/OmicsIntegrator1'), should it be a list?
@@ -65,6 +139,8 @@ def run_container(framework: str, container_suffix: str, command: List[str], vol
         return run_container_docker(container, command, volumes, working_dir, environment)
     elif normalized_framework == 'singularity':
         return run_container_singularity(container, command, volumes, working_dir, environment)
+    elif normalized_framework == 'dsub':
+        return run_container_dsub(container, command, volumes, working_dir, environment)
     else:
         raise ValueError(f'{framework} is not a recognized container framework. Choose "docker" or "singularity".')
 
@@ -258,3 +334,74 @@ def prepare_volume(filename: Union[str, PurePath], volume_base: Union[str, PureP
         src = parent
 
     return (src, dest), container_filename
+
+def run_container_dsub(container: str, command: List[str], volumes: List[Tuple[PurePath, PurePath]], working_dir: str, environment: str = 'SPRAS=True'):
+    """
+    Runs a command in the container using Docker.
+    Attempts to automatically correct file owner and group for new files created by the container, setting them to the
+    current owner and group IDs.
+    Does not modify the owner or group for existing files modified by the container.
+    @param container: name of the container in the Google Cloud Container Registry
+    @param command: command to run in the container
+    @param volumes: a list of volumes to mount where each item is a (source, destination) tuple
+    @param working_dir: the working directory in the container
+    @param environment: environment variables to set in the container
+    @return: path of output from dsub
+    """
+    # Dictionary of flags for dsub command
+    flags = dict()
+
+    workspace_bucket = os.getenv('WORKSPACE_BUCKET')
+    # Add path in the workspace bucket and label for dsub command for each volume
+    dsub_volumes = [(src, dst, workspace_bucket +  str(dst), "INPUT_" + str(i),) for i, (src, dst) in enumerate(volumes)]
+
+    # Prepare command that will be run inside the container for dsub
+    container_command = list()
+    for item in command:
+        # Find if item is volume
+        to_replace = [(str(path[1]), "${"+path[3]+'}') for path in dsub_volumes if str(path[1]) in item]
+        # Replace volume path with dsub volume path
+        if len(to_replace) == 1:
+            # Get path that will be replaced
+            path = to_replace[0][0]
+            # Get dsub input variable that will replace path
+            env_variable = to_replace[0][1]
+            # Replace path with env_variable
+            container_path = item.replace(path, env_variable)
+            # Add / if there is no suffix
+            if container_path == env_variable:
+                container_path = container_path + '/'
+            container_command.append(container_path)
+        else:
+            container_command.append(item)
+
+    # Add a command to copy the volumes to the workspace buckets
+    container_command.append(('; cp -rf ' + f'/mnt/data/input/gs/{workspace_bucket}{working_dir}/*' + ' $OUTPUT').replace('gs://', ''))
+
+    # Make the command into a string
+    flags['command'] = ' '.join(container_command)
+    flags['command'] = "'" + flags['command'] + "'"
+
+    ## Push volumes to WORKSPACE_BUCKET
+    for src, _dst, gcs_path, _env in dsub_volumes:
+        upload_gcs(local_path=str(src), gcs_path=gcs_path, is_dir=True)
+
+    ## Prepare flags for dsub command
+    flags['image'] = container
+    flags['env'] = environment
+    flags['input-recursive'] = [vol[3]+'='+vol[2] for vol in dsub_volumes]
+    flags['output-recursive'] = "OUTPUT=" + workspace_bucket +  working_dir
+    flags['logging'] = workspace_bucket + '/dsub/'
+
+    # Create dsub command
+    dsub_command = prepare_dsub_cmd(flags)
+
+    # Run dsub as subprocess
+    subprocess.run(dsub_command, shell=True)
+
+    # Pull output volumes from WORKSPACE_BUCKET
+    for src, _dst, gcs_path, _env in dsub_volumes:
+        download_gcs(local_path=str(src), gcs_path=gcs_path, is_dir=True)
+
+    # return location of dsub logs in WORKSPACE_BUCKET
+    return  'dsub logs: {logs}'.format(logs = flags['logging'])

spras/meo.py

@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 
 from spras.containers import prepare_volume, run_container
@@ -15,7 +16,7 @@
 # Does not support MINSAT or MAXCSP
 # TODO add parameter validation
 def write_properties(filename=Path('properties.txt'), edges=None, sources=None, targets=None, edge_output=None,
-                     path_output=None, max_path_length=None, local_search=None, rand_restarts=None):
+                     path_output=None, max_path_length=None, local_search=None, rand_restarts=None, framework='docker'):
     """
     Write the properties file for Maximum Edge Orientation
     See https://github.com/agitter/meo/blob/master/sample.props for property descriptions and the default values at
@@ -27,6 +28,17 @@ def write_properties(filename=Path('properties.txt'), edges=None, sources=None,
     if edges is None or sources is None or targets is None or edge_output is None or path_output is None:
         raise ValueError('Required Maximum Edge Orientation properties file arguments are missing')
 
+    if framework == 'dsub':
+        # Get path inside dsub container
+        workspace_bucket = os.getenv('WORKSPACE_BUCKET')
+        input_prefix = f'/mnt/data/input/gs/{workspace_bucket}'.replace('gs://', '')
+        # Add input prefix to all MEO paths
+        edges = input_prefix + edges
+        sources = input_prefix + sources
+        targets = input_prefix + targets
+        edge_output = input_prefix + edge_output
+        path_output = input_prefix + path_output
+
     with open(filename, 'w') as f:
         # Write the required properties
         f.write(f'edges.file = {Path(edges).as_posix()}\n')
@@ -148,7 +160,7 @@ def run(edges=None, sources=None, targets=None, output_file=None, max_path_lengt
         properties_file_local = Path(out_dir, properties_file)
         write_properties(filename=properties_file_local, edges=edge_file, sources=source_file, targets=target_file,
                          edge_output=mapped_output_file, path_output=mapped_path_output,
-                         max_path_length=max_path_length, local_search=local_search, rand_restarts=rand_restarts)
+                         max_path_length=max_path_length, local_search=local_search, rand_restarts=rand_restarts, framework=container_framework)
         bind_path, properties_file = prepare_volume(str(properties_file_local), work_dir)
         volumes.append(bind_path)