dsgrid-toolkit 0.2.0__py3-none-any.whl

dsgrid/cloud/cloud_storage_interface.py

@@ -0,0 +1,140 @@
+import abc
+from contextlib import contextmanager
+
+
+class CloudStorageInterface(abc.ABC):
+    """Defines interface to synchronize data stored on a cloud storage system."""
+
+    @abc.abstractmethod
+    def check_lock_file(self, path):
+        """Checks if a given lock file path exists and that it was created by the same username and uuid.
+
+        Returns an error if the existing lock file's username and uuid do not match.
+
+        Parameters
+        ----------
+        path : str
+            Lock file path
+
+        Raises
+        ------
+        DSGRegistryLockError
+            exception is raised if any lock files are found
+        """
+
+    @abc.abstractmethod
+    def check_valid_lock_file(self, path):
+        """Checks if a given lock file path is valid. Returns errors if invalid.
+
+        Parameters
+        ----------
+        path : str
+            Lock file path
+        """
+
+    @abc.abstractmethod
+    def get_lock_files(self, relative_path=None):
+        """Returns a generator of lock files within the /.locks directory (non-recursive).
+
+        Parameters
+        ----------
+        relative_path : str
+            Relative path to search for lock files in. By default, None.
+        """
+
+    @abc.abstractmethod
+    def has_lock_files(self):
+        """Returns True if a .lock file exists within the /.locks directory."""
+
+    @abc.abstractmethod
+    @contextmanager
+    def make_lock_file_managed(self, path):
+        """Context manager to make a lock file given the file path. On close, it removes the lock file.
+
+        Parameters
+        ----------
+        path : str
+            Lock file path
+
+        Raises
+        ------
+        DSGRegistryLockError
+            Raises a registry lock error if a registry.lock already exists
+        """
+
+    @abc.abstractmethod
+    def make_lock_file(self, path):
+        """Make a lock file given the file path.
+
+        Parameters
+        ----------
+        path : str
+            Lock file path
+
+        Raises
+        ------
+        DSGRegistryLockError
+            Raises a registry lock error if a registry.lock already exists
+        """
+
+    @abc.abstractmethod
+    def read_lock_file(self, path):
+        """Reads a lock file and returns a dictionary of its contents.
+
+        Parameters
+        ----------
+        path : str
+            Lock file path
+        """
+
+    @abc.abstractmethod
+    def remove_lock_file(self, path, force=False):
+        """Remove a lock file.
+
+        Parameters
+        ----------
+        path : str
+            Lock file path
+        force : bool
+            Boolean flag to force removal of lock file that does not have the same UUID or username, by default False
+
+        Raises
+        ------
+        DSGRegistryLockError
+            Raises a registry lock error if a registry.lock already exists and is was generated by a different user or different UUID (if force==False).
+        """
+
+    @abc.abstractmethod
+    def sync_pull(self, remote_path, local_path, exclude=None, delete_local=False, is_file=False):
+        """Synchronize data from remote_path to local_path.
+        If delete_local is True, this deletes any files in local_path that do not exist in remote_path.
+
+        Parameters
+        ----------
+        remote_path : str
+            Remote registry path
+        local_path : str
+            Local registry path
+        delete_local : bool, optional
+            If true, this method deletes files and directories that exist in the local_path but not in the remote_path
+        exclude : list, optional
+            List of patterns to exclude, by default None.
+            If excluding whole directories, the exclusion must end with /* , e.g. data/*
+        is_file : bool, optional
+            If the path is a file (not a registry). By default False.
+        """
+
+    @abc.abstractmethod
+    def sync_push(self, remote_path, local_path, exclude=None):
+        """Synchronize data from local path to remote_path
+
+        Parameters
+        ----------
+        remote_path : str
+            Remote registry path
+        local_path : str
+            Local registry path
+        exclude : list, optional
+            List of patterns to exclude, by default None.
+            If excluding whole directories, the exclusion must end with /* , e.g. data/*
+        """

dsgrid/cloud/factory.py

@@ -0,0 +1,31 @@
+# from .s3_storage_interface import S3StorageInterface
+from .fake_storage_interface import FakeStorageInterface
+
+# from dsgrid.common import AWS_PROFILE_NAME
+
+
+def make_cloud_storage_interface(local_path, remote_path, uuid, user, offline=False):
+    """Creates a CloudStorageInterface appropriate for path.
+
+    Parameters
+    ----------
+    local_path : str
+    remote_path : str
+    uuid : str
+        Unique ID to be used when generating cloud locks.
+    user : str
+        Username to be used when generating cloud locks.
+    offline : bool, optional
+        If True, don't perform any remote syncing operations.
+
+
+    Returns
+    -------
+    CloudStorageInterface
+
+    """
+    if not offline and remote_path.lower().startswith("s3"):
+        msg = f"Support for S3 is currently disabled: {remote_path=}"
+        raise NotImplementedError(msg)
+        # return S3StorageInterface(local_path, remote_path, uuid, user, profile=AWS_PROFILE_NAME)
+    return FakeStorageInterface()

dsgrid/cloud/fake_storage_interface.py

@@ -0,0 +1,37 @@
+from contextlib import contextmanager
+from .cloud_storage_interface import CloudStorageInterface
+
+
+class FakeStorageInterface(CloudStorageInterface):
+    """Fake interface for tests and local mode."""
+
+    def check_lock_file(self, path):
+        pass
+
+    def check_valid_lock_file(self, path):
+        pass
+
+    def get_lock_files(self, relative_path=None):
+        pass
+
+    def has_lock_files(self):
+        pass
+
+    @contextmanager
+    def make_lock_file_managed(self, path):
+        yield
+
+    def make_lock_file(self, path):
+        pass
+
+    def read_lock_file(self, path):
+        pass
+
+    def remove_lock_file(self, path, force=False):
+        pass
+
+    def sync_pull(self, remote_path, local_path, exclude=None, delete_local=False, is_file=False):
+        pass
+
+    def sync_push(self, remote_path, local_path, exclude=None):
+        pass

dsgrid/cloud/s3_storage_interface.py

@@ -0,0 +1,156 @@
+from contextlib import contextmanager
+from datetime import datetime
+import json
+import logging
+import os
+import sys
+from pathlib import Path
+import time
+
+from dsgrid.exceptions import DSGMakeLockError, DSGRegistryLockError
+from dsgrid.filesystem.local_filesystem import LocalFilesystem
+from dsgrid.filesystem.s3_filesystem import S3Filesystem
+from dsgrid.utils.run_command import check_run_command
+from dsgrid.utils.timing import track_timing, timer_stats_collector
+from .cloud_storage_interface import CloudStorageInterface
+
+logger = logging.getLogger(__name__)
+
+
+class S3StorageInterface(CloudStorageInterface):
+    """Interface to S3."""
+
+    def __init__(self, local_path, remote_path, uuid, user, profile):
+        self._local_path = local_path
+        self._remote_path = remote_path
+        self._uuid = uuid
+        self._user = user
+        self._local_filesystem = LocalFilesystem()
+        self._s3_filesystem = S3Filesystem(remote_path, profile)
+
+    def _sync(self, src, dst, exclude=None, is_file=False):
+        start = time.time()
+        aws_exec = self._get_aws_executable()
+        cmd = "cp" if is_file else "sync"
+        sync_command = f"{aws_exec} s3 {cmd} {src} {dst} --profile {self._s3_filesystem.profile}"
+        if exclude:
+            for x in exclude:
+                sync_command = sync_command + f" --exclude {x}"
+        logger.info("Running %s", sync_command)
+        check_run_command(sync_command)
+        logger.info("Command took %s seconds", time.time() - start)
+
+    @staticmethod
+    def _get_aws_executable():
+        # subprocess.run cannot find the aws executable on Windows if shell is not True.
+        # That's probably an aws cli bug. We can workaround it here.
+        return "aws.cmd" if sys.platform == "win32" else "aws"
+
+    def check_lock_file(self, path):
+        self.check_valid_lock_file(path)
+        filepath = self._s3_filesystem.path(path)
+        if filepath.exists():
+            lock_contents = self.read_lock_file(filepath)
+            if (
+                not self._uuid == lock_contents["uuid"]
+                or not self._user == lock_contents["username"]
+            ):
+                msg = f"Registry path {str(filepath)} is currently locked by {lock_contents['username']}, timestamp={lock_contents['timestamp']}, uuid={lock_contents['uuid']}."
+                raise DSGRegistryLockError(msg)
+
+    def check_valid_lock_file(self, path):
+        path = Path(path)
+        # check that lock file is of type .lock
+        if path.suffix != ".lock":
+            msg = f"Lock file path provided ({path}) must be a valid .lock path"
+            raise DSGMakeLockError(msg)
+        # check that lock file in expected dirs
+        relative_path = Path(path).parent
+        if str(relative_path).startswith("s3:/nrel-dsgrid-registry/"):
+            relative_path = relative_path.relative_to("s3:/nrel-dsgrid-registry/")
+        if str(relative_path).startswith("/"):
+            relative_path = relative_path.relative_to("/")
+        if relative_path == Path("configs/.locks"):
+            pass
+        elif relative_path == Path("data/.locks"):
+            pass
+        else:
+            DSGMakeLockError(
+                "Lock file path provided must have relative path of configs/.locks or data/.locks"
+            )
+        return True
+
+    def get_lock_files(self, relative_path=None):
+        if relative_path:
+            contents = self._s3_filesystem.path(
+                f"{self._s3_filesystem.bucket}/{relative_path}"
+            ).glob(pattern="**/*.locks")
+        else:
+            contents = self._s3_filesystem.path(self._s3_filesystem.bucket).glob(
+                pattern="**/*.locks"
+            )
+        return contents
+
+    def has_lock_files(self):
+        contents = self.get_lock_files()
+        return next(contents, None) is not None
+
+    @contextmanager
+    def make_lock_file_managed(self, path):
+        try:
+            self.make_lock_file(path)
+            yield
+        finally:
+            self.remove_lock_file(path)
+
+    def make_lock_file(self, path):
+        self.check_lock_file(path)
+        filepath = self._s3_filesystem.path(path)
+        lock_content = {
+            "username": self._user,
+            "uuid": self._uuid,
+            "timestamp": str(datetime.now()),
+        }
+        self._s3_filesystem.path(filepath).write_text(json.dumps(lock_content))
+
+    def read_lock_file(self, path):
+        lockfile_contents = json.loads(self._s3_filesystem.path(path).read_text())
+        return lockfile_contents
+
+    def remove_lock_file(self, path, force=False):
+        filepath = self._s3_filesystem.path(path)
+        if filepath.exists():
+            lockfile_contents = self.read_lock_file(filepath)
+            if not force:
+                if (
+                    not self._uuid == lockfile_contents["uuid"]
+                    and not self._user == lockfile_contents["username"]
+                ):
+                    msg = f"Registry path {str(filepath)} is currently locked by {lockfile_contents['username']}. Lock created as {lockfile_contents['timestamp']} with uuid={lockfile_contents['uuid']}."
+                    raise DSGRegistryLockError(msg)
+            if force:
+                logger.warning(
+                    "Force removed lock file with user=%s and uuid=%s",
+                    lockfile_contents["username"],
+                    lockfile_contents["uuid"],
+                )
+            filepath.unlink()
+
+    @track_timing(timer_stats_collector)
+    def sync_pull(self, remote_path, local_path, exclude=None, delete_local=False, is_file=False):
+        if delete_local:
+            local_contents = self._local_filesystem.rglob(local_path)
+            s3_contents = {
+                str(self._s3_filesystem.path(x).relative_to(self._s3_filesystem.path(remote_path)))
+                for x in self._s3_filesystem.rglob(remote_path)
+            }
+            for content in local_contents:
+                relcontent = os.path.relpath(content, local_path)
+                if relcontent not in s3_contents:
+                    self._local_filesystem.rm(content)
+                    logger.info("delete: %s because it is not in %s", relcontent, remote_path)
+        self._sync(remote_path, local_path, exclude, is_file)
+
+    @track_timing(timer_stats_collector)
+    def sync_push(self, remote_path, local_path, exclude=None):
+        self._sync(local_path, remote_path, exclude)

dsgrid/common.py

@@ -0,0 +1,35 @@
+import enum
+import os
+from pathlib import Path
+
+# AWS_PROFILE_NAME = "nrel-aws-dsgrid"
+REMOTE_REGISTRY = "s3://nrel-dsgrid-registry"
+
+
+def on_hpc():
+    # NREL_CLUSTER is not set when you ssh into a compute node.
+    return "NREL_CLUSTER" in os.environ or "SLURM_JOB_ID" in os.environ
+
+
+if on_hpc():
+    LOCAL_REGISTRY = Path("/scratch") / os.environ["USER"] / ".dsgrid-registry"
+else:
+    LOCAL_REGISTRY = Path.home() / ".dsgrid-registry"
+
+LOCAL_REGISTRY_DATA = LOCAL_REGISTRY / "data"
+PROJECT_FILENAME = "project.json5"
+REGISTRY_FILENAME = "registry.json5"
+DATASET_FILENAME = "dataset.json5"
+DIMENSIONS_FILENAME = "dimensions.json5"
+DEFAULT_DB_PASSWORD = "openSesame"
+DEFAULT_SCRATCH_DIR = "__dsgrid_scratch__"
+SCALING_FACTOR_COLUMN = "scaling_factor"
+SYNC_EXCLUDE_LIST = ["*.DS_Store", "**/*.lock"]
+VALUE_COLUMN = "value"
+
+
+class BackendEngine(enum.StrEnum):
+    """Supported backend SQL processing engines"""
+
+    DUCKDB = "duckdb"
+    SPARK = "spark"

dsgrid/config/annual_time_dimension_config.py

@@ -0,0 +1,187 @@
+import logging
+from datetime import timedelta, datetime
+
+import pandas as pd
+from chronify.time_range_generator_factory import make_time_range_generator
+
+from dsgrid.config.date_time_dimension_config import DateTimeDimensionConfig
+from dsgrid.dimension.base_models import DimensionType
+from dsgrid.dimension.time import AnnualTimeRange
+from dsgrid.exceptions import DSGInvalidDataset
+from dsgrid.time.types import AnnualTimestampType
+from dsgrid.dimension.time_utils import is_leap_year
+from dsgrid.spark.functions import (
+    cross_join,
+    handle_column_spaces,
+    select_expr,
+)
+from dsgrid.spark.types import (
+    DataFrame,
+    StructType,
+    StructField,
+    IntegerType,
+    StringType,
+    TimestampType,
+    F,
+)
+from dsgrid.utils.timing import timer_stats_collector, track_timing
+from dsgrid.utils.spark import (
+    get_spark_session,
+    set_session_time_zone,
+)
+from .dimensions import AnnualTimeDimensionModel
+from .time_dimension_base_config import TimeDimensionBaseConfig
+
+
+logger = logging.getLogger(__name__)
+
+
+class AnnualTimeDimensionConfig(TimeDimensionBaseConfig):
+    """Provides an interface to an AnnualTimeDimensionModel."""
+
+    @staticmethod
+    def model_class() -> AnnualTimeDimensionModel:
+        return AnnualTimeDimensionModel
+
+    @track_timing(timer_stats_collector)
+    def check_dataset_time_consistency(self, load_data_df, time_columns) -> None:
+        logger.info("Check AnnualTimeDimensionConfig dataset time consistency.")
+        if len(time_columns) > 1:
+            msg = (
+                "AnnualTimeDimensionConfig expects only one column from "
+                f"get_load_data_time_columns, but has {time_columns}"
+            )
+            raise ValueError(msg)
+        time_col = time_columns[0]
+        time_ranges = self.get_time_ranges()
+        assert len(time_ranges) == 1, len(time_ranges)
+        time_range = time_ranges[0]
+        # TODO: need to support validation of multiple time ranges: DSGRID-173
+
+        expected_timestamps = time_range.list_time_range()
+        actual_timestamps = [
+            pd.Timestamp(str(x[time_col]), tz=self.get_tzinfo()).to_pydatetime()
+            for x in load_data_df.select(time_col)
+            .distinct()
+            .filter(f"{time_col} IS NOT NULL")
+            .sort(time_col)
+            .collect()
+        ]
+        if expected_timestamps != actual_timestamps:
+            mismatch = sorted(
+                set(expected_timestamps).symmetric_difference(set(actual_timestamps))
+            )
+            msg = f"load_data {time_col}s do not match expected times. mismatch={mismatch}"
+            raise DSGInvalidDataset(msg)
+
+    def build_time_dataframe(self) -> DataFrame:
+        time_col = self.get_load_data_time_columns()
+        assert len(time_col) == 1, time_col
+        time_col = time_col[0]
+        schema = StructType([StructField(time_col, IntegerType(), False)])
+
+        model_time = self.list_expected_dataset_timestamps()
+        df_time = get_spark_session().createDataFrame(model_time, schema=schema)
+        return df_time
+
+    def get_frequency(self) -> timedelta:
+        return timedelta(days=365)
+
+    def get_time_ranges(self) -> list[AnnualTimeRange]:
+        ranges = []
+        frequency = self.get_frequency()
+        for start, end in self._build_time_ranges(
+            self.model.ranges, self.model.str_format, tz=self.get_tzinfo()
+        ):
+            start = pd.Timestamp(start)
+            end = pd.Timestamp(end)
+            ranges.append(
+                AnnualTimeRange(
+                    start=start,
+                    end=end,
+                    frequency=frequency,
+                )
+            )
+
+        return ranges
+
+    def get_start_times(self) -> list[pd.Timestamp]:
+        tz = self.get_tzinfo()
+        start_times = []
+        for trange in self.model.ranges:
+            start = datetime.strptime(trange.start, self.model.str_format)
+            assert start.tzinfo is None
+            start_times.append(start.replace(tzinfo=tz))
+
+        return start_times
+
+    def get_lengths(self) -> list[int]:
+        lengths = []
+        for trange in self.model.ranges:
+            start = datetime.strptime(trange.start, self.model.str_format)
+            end = datetime.strptime(trange.end, self.model.str_format)
+            lengths.append(end.year - start.year + 1)
+        return lengths
+
+    def get_load_data_time_columns(self) -> list[str]:
+        return list(AnnualTimestampType._fields)
+
+    def get_time_zone(self) -> None:
+        return None
+
+    def get_tzinfo(self) -> None:
+        return None
+
+    def get_time_interval_type(self) -> None:
+        return None
+
+    def list_expected_dataset_timestamps(self) -> list[AnnualTimestampType]:
+        timestamps = []
+        for time_range in self.model.ranges:
+            start, end = (int(time_range.start), int(time_range.end))
+            timestamps += [AnnualTimestampType(x) for x in range(start, end + 1)]
+        return timestamps
+
+
+def map_annual_time_to_date_time(
+    df: DataFrame,
+    annual_dim: AnnualTimeDimensionConfig,
+    dt_dim: DateTimeDimensionConfig,
+    value_columns: set[str],
+) -> DataFrame:
+    """Map a DataFrame with an annual time dimension to a DateTime time dimension."""
+    annual_col = annual_dim.get_load_data_time_columns()[0]
+    myear_column = DimensionType.MODEL_YEAR.value
+    timestamps = make_time_range_generator(dt_dim.to_chronify()).list_timestamps()
+    time_cols = dt_dim.get_load_data_time_columns()
+    assert len(time_cols) == 1, time_cols
+    time_col = time_cols[0]
+    schema = StructType([StructField(time_col, TimestampType(), False)])
+    dt_df = get_spark_session().createDataFrame(
+        [(x.to_pydatetime(),) for x in timestamps], schema=schema
+    )
+
+    # Note that MeasurementType.TOTAL has already been verified.
+    with set_session_time_zone(dt_dim.model.datetime_format.timezone.tz_name):
+        years = (
+            select_expr(dt_df, [f"YEAR({handle_column_spaces(time_col)}) AS year"])
+            .distinct()
+            .collect()
+        )
+        if len(years) != 1:
+            msg = "DateTime dimension has more than one year: {years=}"
+            raise NotImplementedError(msg)
+        if annual_dim.model.include_leap_day and is_leap_year(years[0].year):
+            measured_duration = timedelta(days=366)
+        else:
+            measured_duration = timedelta(days=365)
+
+    df2 = (
+        cross_join(df, dt_df)
+        .withColumn(myear_column, F.col(annual_col).cast(StringType()))
+        .drop(annual_col)
+    )
+    frequency = dt_dim.model.frequency
+    for column in value_columns:
+        df2 = df2.withColumn(column, F.col(column) / (measured_duration / frequency))
+    return df2