Merge pull request #241 from mila-iqia/refactor

Plugin refactor for diskusage.

Author: satyaog

config/sarc-dev.yaml

@@ -27,8 +27,15 @@ sarc:
       sacct_bin: "/opt/slurm/bin/sacct"
       duc_inodes_command:
       duc_storage_command:
-      diskusage_report_command: beegfs-ctl --cfgFile=/etc/beegfs/home.d/beegfs-client.conf
-        --getquota --uid $USER --csv
+      diskusage:
+         - name: beegfs
+           params:
+             config_files:
+               default: /etc/beegfs/scratch.d/beegfs-client.conf
+         - name: beegfs
+           params:
+             config_files:
+               default: /etc/beegfs/projects.d/beegfs-client.conf
       prometheus_url: http://prometheus01.server.mila.quebec:9090/
       start_date: '2022-04-01'
       billing_is_gpu: true
@@ -81,7 +88,8 @@ sarc:
         --count /project/rrg-bengioy-ad
       duc_storage_command: duc ls -d /project/.duc_databases/rrg-bengioy-ad.sqlite
         /project/rrg-bengioy-ad
-      diskusage_report_command: diskusage_report --project --all_users
+      diskusage:
+        - name: drac
       prometheus_url: https://mila-thanos.calculquebec.ca
       prometheus_headers_file: ../../SARC_secrets/secrets/drac_prometheus/headers.json
       start_date: '2022-04-01'

sarc/cli/acquire/storages.py

@@ -1,21 +1,14 @@
 import logging
 from dataclasses import dataclass
-from typing import Callable
 
 from simple_parsing import field
 
-from sarc.config import ClusterConfig, config
-from sarc.storage.diskusage import DiskUsage, get_diskusage_collection
-from sarc.storage.drac import fetch_diskusage_report as fetch_dirac_diskusage
-from sarc.storage.mila import fetch_diskusage_report as fetch_mila_diskusage
+from sarc.config import config
+from sarc.core.scraping.diskusage import get_diskusage_scraper
+from sarc.storage.diskusage import get_diskusage_collection
 
 logger = logging.getLogger(__name__)
 
-methods: dict[str, Callable[[ClusterConfig], DiskUsage]] = {
-    "default": fetch_dirac_diskusage,
-    "mila": fetch_mila_diskusage,
-}
-
 
 @dataclass
 class AcquireStorages:
@@ -29,14 +22,30 @@ def execute(self) -> int:
             logger.info(f"Acquiring {cluster_name} storages report...")
 
             cluster = cfg.clusters[cluster_name]
-
-            fetch_diskusage = methods.get(cluster_name, methods["default"])
-            du = fetch_diskusage(cluster)
-
-            if not self.dry:
-                collection = get_diskusage_collection()
-                collection.add(du)
-            else:
-                logger.info("Document:\n" + du.model_dump_json(indent=2))
+            diskusage_configs = cluster.diskusage
+            if diskusage_configs is None:
+                continue
+
+            # Process each diskusage configuration
+            for diskusage_config in diskusage_configs:
+                try:
+                    scraper = get_diskusage_scraper(diskusage_config.name)
+                except KeyError as ke:
+                    logger.exception(
+                        "Invalid or absent diskusage scraper name: %s",
+                        diskusage_config.name,
+                        exc_info=ke,
+                    )
+                    continue
+
+                disk_config = scraper.validate_config(diskusage_config.params)
+                data = scraper.get_diskusage_report(cluster.ssh, disk_config)
+                du = scraper.parse_diskusage_report(disk_config, cluster_name, data)
+
+                if not self.dry:
+                    collection = get_diskusage_collection()
+                    collection.add(du)
+                else:
+                    logger.info("Document:\n" + du.model_dump_json(indent=2))
 
         return 0

sarc/config.py

@@ -17,6 +17,8 @@
 
 from .alerts.common import HealthMonitorConfig
 
+type JSON = list[JSON] | dict[str, JSON] | int | str | float | bool | None
+
 if TYPE_CHECKING:
     from fabric import Connection
     from prometheus_api_client import PrometheusConnect
@@ -37,6 +39,12 @@ class ConfigurationError(Exception):
     pass
 
 
+@dataclass
+class DiskUsageConfig:
+    name: str
+    params: JSON = field(default_factory=dict)
+
+
 @dataclass
 class ClusterConfig:
     # pylint: disable=too-many-instance-attributes
@@ -51,7 +59,7 @@ class ClusterConfig:
     sshconfig: Path | None = None
     duc_inodes_command: str | None = None
     duc_storage_command: str | None = None
-    diskusage_report_command: str | None = None
+    diskusage: list[DiskUsageConfig] | None = None
     start_date: str = "2022-04-01"
     slurm_conf_host_path: Path = Path("/etc/slurm/slurm.conf")
 

sarc/core/models/diskusage.py

@@ -0,0 +1,24 @@
+from pydantic import BaseModel, ByteSize
+
+from .validators import datetime_utc
+
+
+class DiskUsageUser(BaseModel):
+    user: str
+    nbr_files: int
+    size: ByteSize
+
+
+class DiskUsageGroup(BaseModel):
+    group_name: str
+    users: list[DiskUsageUser]
+
+
+class DiskUsage(BaseModel):
+    """
+    Disk usage on a given cluster
+    """
+
+    cluster_name: str
+    groups: list[DiskUsageGroup]
+    timestamp: datetime_utc

sarc/core/models/validators.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from typing import Annotated, Any, Callable
+
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import CoreSchema, core_schema
+
+
+@dataclass(frozen=True)
+class DatetimeUTCValidator:
+    def validate_tz_utc(self, value: datetime, handler: Callable):
+        assert value.tzinfo == UTC, "date is not in UTC time"
+
+        return handler(value)
+
+    def __get_pydantic_core_schema__(
+        self, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        return core_schema.no_info_wrap_validator_function(
+            self.validate_tz_utc, handler(source_type)
+        )
+
+
+type datetime_utc = Annotated[datetime, DatetimeUTCValidator()]

sarc/core/scraping/diskusage.py

@@ -0,0 +1,39 @@
+from importlib.metadata import entry_points
+from typing import Any, Protocol, Type
+
+from fabric import Connection
+from serieux import deserialize
+
+from sarc.core.models.diskusage import DiskUsage
+
+
+class DiskUsageScraper[T](Protocol):
+    config_type: Type[T]
+
+    def validate_config(self, config_data: Any) -> T:
+        """Validate the configuration"""
+        return deserialize(self.config_type, config_data)
+
+    def get_diskusage_report(self, ssh: Connection, config: T) -> bytes:
+        """Get the raw disk usage report for caching and archiving purposes"""
+        ...
+
+    def parse_diskusage_report(
+        self, config: T, cluster_name: str, data: bytes
+    ) -> DiskUsage:
+        """Parse previously fetched report into a DiskUsage"""
+        ...
+
+
+_builtin_scrapers: dict[str, DiskUsageScraper] = dict()
+_diskusage_scrapers = entry_points(group="sarc.diskusage_scraper")
+
+
+def get_diskusage_scraper(name: str) -> DiskUsageScraper:
+    """Raises KeyError if the name is not found"""
+    try:
+        return _builtin_scrapers[name]
+    except KeyError:
+        pass
+    val = _diskusage_scrapers[name]
+    return val.load()

sarc/core/utils.py

@@ -0,0 +1,18 @@
+from fabric import Connection
+
+
+def run_command(
+    connection: Connection, command: str, retries: int
+) -> tuple[str | None, list[Exception]]:
+    errors: list[Exception] = []
+
+    for _ in range(retries):
+        try:
+            result = connection.run(command, hide=True)
+            return result.stdout, errors
+
+        # pylint: disable=broad-exception-caught
+        except Exception as err:
+            errors.append(err)
+
+    return None, errors

sarc/storage/__init__.py

@@ -0,0 +1 @@
+from . import beegfs, drac  # noqa: F401 (these import are to register the modules)

sarc/storage/beegfs.py

@@ -0,0 +1,107 @@
+"""
+Fetching and parsing code specific to the mila cluster
+"""
+
+import csv
+import json
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from io import StringIO
+from typing import cast
+
+from fabric import Connection
+from pydantic import ByteSize
+
+from sarc.client.users.api import User, get_users
+from sarc.core.models.diskusage import DiskUsage, DiskUsageGroup, DiskUsageUser
+from sarc.core.scraping.diskusage import DiskUsageScraper, _builtin_scrapers
+from sarc.core.utils import run_command
+
+beegfs_header = "name,id,size,hard,files,hard"
+
+
+@dataclass
+class BeeGFSDiskUsageConfig:
+    config_files: dict[str, str]
+    retries: int = 3
+    beegfs_ctl_path: str = "beegfs-ctl"
+
+
+class BeeGFSDiskUsage(DiskUsageScraper[BeeGFSDiskUsageConfig]):
+    config_type = BeeGFSDiskUsageConfig
+
+    def get_diskusage_report(
+        self, ssh: Connection, config: BeeGFSDiskUsageConfig
+    ) -> bytes:
+        users = get_users()
+        assert len(users) > 0
+
+        errors: list[Exception] = []
+        failures: list[User] = []
+        output: dict[str, list[str]] = {}
+
+        for name, file in config.config_files.items():
+            usage: list[str] = []
+
+            for user in users:
+                if not user.mila.active:
+                    continue
+
+                cmd = f"{config.beegfs_ctl_path} --cfgFile={file} --getquota --uid {user.mila.username} --csv"
+                result, err = run_command(ssh, cmd, config.retries)
+
+                if err:
+                    errors.extend(err)
+
+                if result is None:
+                    failures.append(user)
+                else:
+                    usage.append(_trim_beegfs_output(result))
+
+            output[name] = usage
+
+        return json.dumps(output).encode()
+
+    def parse_diskusage_report(
+        self, config: BeeGFSDiskUsageConfig, cluster_name: str, data_raw: bytes
+    ) -> DiskUsage:
+        groups: list[DiskUsageGroup] = []
+        data: dict[str, str] = json.loads(data_raw.decode())
+
+        groups = [
+            DiskUsageGroup(
+                group_name=name,
+                users=[
+                    _parse_line(line)
+                    for line in csv.reader(StringIO("\n".join(data[name])))
+                ],
+            )
+            for name in config.config_files.keys()
+        ]
+
+        return DiskUsage(
+            cluster_name=cluster_name,
+            groups=groups,
+            timestamp=datetime.now(UTC),
+        )
+
+
+# Register the scraper to make it available
+_builtin_scrapers["beegfs"] = BeeGFSDiskUsage()
+
+
+def _trim_beegfs_output(output: str) -> str:
+    splitted = output.splitlines()
+    header_index = splitted.index(beegfs_header)
+    return "\n".join(splitted[header_index + 1 :])
+
+
+def _parse_line(line: list[str]) -> DiskUsageUser:
+    columns = {
+        key.strip(): value.strip() for key, value in zip(beegfs_header.split(","), line)
+    }
+    return DiskUsageUser(
+        user=columns["name"],
+        nbr_files=int(columns["files"]),
+        size=cast(ByteSize, columns["size"]),
+    )