Merge pull request #409 from rewiringamerica/refactor

Small refactor to simplify adding a GCP implementation

buildstockbatch/aws/aws.py

@@ -16,10 +16,8 @@
 from botocore.exceptions import ClientError
 import collections
 import csv
-import docker
 from fsspec.implementations.local import LocalFileSystem
 import gzip
-import hashlib
 import itertools
 from joblib import Parallel, delayed
 import json
@@ -38,11 +36,13 @@
 import io
 import zipfile
 
-from buildstockbatch.base import ValidationError, BuildStockBatchBase
+from buildstockbatch.base import ValidationError
 from buildstockbatch.aws.awsbase import AwsJobBase
+from buildstockbatch.cloud.docker_base import DockerBatchBase
 from buildstockbatch import postprocessing
 from buildstockbatch.utils import (
-    ContainerRuntime,
+    calc_hash_for_file,
+    compress_file,
     log_error_details,
     get_project_configuration,
     read_csv,
@@ -76,17 +76,6 @@ def filename_generator():
     )
 
 
-def compress_file(in_filename, out_filename):
-    with gzip.open(str(out_filename), "wb") as f_out:
-        with open(str(in_filename), "rb") as f_in:
-            shutil.copyfileobj(f_in, f_out)
-
-
-def calc_hash_for_file(filename):
-    with open(filename, "rb") as f:
-        return hashlib.sha256(f.read()).hexdigest()
-
-
 def copy_s3_file(src_bucket, src_key, dest_bucket, dest_key):
     s3 = boto3.client("s3")
     s3.copy({"Bucket": src_bucket, "Key": src_key}, dest_bucket, dest_key)
@@ -1488,28 +1477,6 @@ def clean(self):
         logger.info(f"Simple notifications topic {self.sns_state_machine_topic} deleted.")
 
 
-class DockerBatchBase(BuildStockBatchBase):
-    CONTAINER_RUNTIME = ContainerRuntime.DOCKER
-
-    def __init__(self, project_filename):
-        super().__init__(project_filename)
-
-        self.docker_client = docker.DockerClient.from_env()
-        try:
-            self.docker_client.ping()
-        except:  # noqa: E722 (allow bare except in this case because error can be a weird non-class Windows API error)
-            logger.error("The docker server did not respond, make sure Docker Desktop is started then retry.")
-            raise RuntimeError("The docker server did not respond, make sure Docker Desktop is started then retry.")
-
-    @staticmethod
-    def validate_project(project_file):
-        super(DockerBatchBase, DockerBatchBase).validate_project(project_file)
-
-    @property
-    def docker_image(self):
-        return "nrel/openstudio:{}".format(self.os_version)
-
-
 class AwsBatch(DockerBatchBase):
     def __init__(self, project_filename):
         super().__init__(project_filename)

buildstockbatch/cloud/docker_base.py

@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+
+"""
+buildstockbatch.docker_base
+~~~~~~~~~~~~~~~
+This is the base class mixed into classes that deploy using a docker container.
+
+:author: Natalie Weires
+:license: BSD-3
+"""
+import docker
+import logging
+
+from buildstockbatch.base import BuildStockBatchBase
+from buildstockbatch.utils import ContainerRuntime
+
+logger = logging.getLogger(__name__)
+
+
+class DockerBatchBase(BuildStockBatchBase):
+    """Base class for implementations that run in Docker containers."""
+
+    CONTAINER_RUNTIME = ContainerRuntime.DOCKER
+
+    def __init__(self, project_filename):
+        super().__init__(project_filename)
+
+        self.docker_client = docker.DockerClient.from_env()
+        try:
+            self.docker_client.ping()
+        except:  # noqa: E722 (allow bare except in this case because error can be a weird non-class Windows API error)
+            logger.error("The docker server did not respond, make sure Docker Desktop is started then retry.")
+            raise RuntimeError("The docker server did not respond, make sure Docker Desktop is started then retry.")
+
+    @staticmethod
+    def validate_project(project_file):
+        super(DockerBatchBase, DockerBatchBase).validate_project(project_file)
+
+    @property
+    def docker_image(self):
+        return "nrel/openstudio:{}".format(self.os_version)
+
+    @property
+    def weather_dir(self):
+        return self._weather_dir

buildstockbatch/utils.py

@@ -1,10 +1,13 @@
 import enum
+import gzip
+import hashlib
 import inspect
-import os
 import logging
+import os
+import pandas as pd
+import shutil
 import traceback
 import yaml
-import pandas as pd
 
 
 logger = logging.getLogger(__name__)
@@ -125,3 +128,14 @@ def run_with_error_capture(*args, **kwargs):
         return run_with_error_capture
 
     return log_error_decorator
+
+
+def compress_file(in_filename, out_filename):
+    with gzip.open(str(out_filename), "wb") as f_out:
+        with open(str(in_filename), "rb") as f_in:
+            shutil.copyfileobj(f_in, f_out)
+
+
+def calc_hash_for_file(filename):
+    with open(filename, "rb") as f:
+        return hashlib.sha256(f.read()).hexdigest()