so-campaign-manager 0.0.4__py3-none-any.whl

socm/planner/heft_planner.py

@@ -0,0 +1,442 @@
+from typing import Dict, List, Tuple
+
+import networkx as nx
+import numpy as np
+
+from ..core import Campaign, QosPolicy, Resource, Workflow
+from .base import PlanEntry, Planner
+
+
+class HeftPlanner(Planner):
+    """Campaign planner using Heterogeneous Earliest Finish Time (HEFT) algorithm.
+
+    HEFT is a list scheduling algorithm that assigns tasks to processors to minimize
+    the overall completion time, considering both computation and communication costs.
+
+    Reference:
+        Topcuoglu, H., Hariri, S., & Wu, M. Y. (2002). Performance-effective and
+        low-complexity task scheduling for heterogeneous computing.
+        IEEE Transactions on Parallel and Distributed Systems, 13(3), 260-274.
+
+    Attributes:
+        _estimated_walltime: List of estimated execution times (walltime) for each workflow
+        _estimated_cpus: List of estimated CPU requirements for each workflow
+        _estimated_memory: List of estimated memory requirements for each workflow
+    """
+
+    def __init__(
+        self,
+        campaign: Campaign | None = None,
+        resources: Resource | None = None,
+        resource_requirements: Dict[int, Dict[str, float]] | None = None,
+        policy: str | None = None,
+        sid: str | None = None,
+        objective: int | None = None
+    ):
+        super().__init__(
+            campaign=campaign,
+            resources=resources,
+            resource_requirements=resource_requirements,
+            sid=sid,
+            policy=policy,
+            objective=objective
+        )
+        # Initialize estimation tables (populated during planning)
+        self._estimated_walltime: List[float] = []
+        self._estimated_cpus: List[int] = []
+        self._estimated_memory: List[float] = []
+
+    def _get_free_memory(self, start_time: float, num_nodes: float) -> float:
+        """Calculate available memory at a given time.
+
+        Args:
+            start_time: Time point to check memory availability
+            num_nodes: The total number of nodes used
+
+        Returns:
+            Available memory in MB
+        """
+        total_memory = num_nodes * self._resources.memory_per_node
+        used_memory = sum(
+            entry.memory
+            for entry in self._plan
+            if entry.start_time <= start_time < entry.end_time
+        )
+        return total_memory - used_memory
+
+    def _get_max_ncores(self, resource_requirements: Dict[int, Dict[str, float]]) -> int:
+        """Get the maximum number of cores required by any single workflow."""
+        return max(values["req_cpus"] for values in resource_requirements.values())
+
+    def _find_suitable_qos_policies(self, requested_cores: int) -> QosPolicy:
+        """Find QoS policies that can accommodate the campaign deadline."""
+        suitable_qos = self._resources.fits_in_qos(self._objective, cores=requested_cores)
+        if not suitable_qos:
+            available_qos = ', '.join(f"{q.name}({q.max_walltime}min)" for q in self._resources.qos)
+            raise ValueError(
+                f"No QoS policy can accommodate deadline of {self._objective} minutes. "
+                f"Available policies: {available_qos}"
+            )
+
+        return suitable_qos
+
+    def _binary_search_resources(
+        self,
+        campaign: List[Workflow],
+        resource_requirements: Dict[int, Dict[str, float]],
+        lower_bound: int,
+        upper_bound: int
+    ) -> Tuple[int, List[PlanEntry], nx.DiGraph] | None:
+        """Binary search for minimum resources that satisfy the deadline.
+
+        Returns:
+            Tuple of (ncores, plan, graph) if successful, None otherwise.
+        """
+        best_ncores = None
+        best_plan = None
+        best_graph = None
+
+        while lower_bound <= upper_bound:
+            mid = (lower_bound + upper_bound) // 2
+
+            test_plan, test_graph = self._calculate_plan(
+                campaign=campaign,
+                resources=range(mid),
+                resource_requirements=resource_requirements
+            )
+
+            max_finish_time = max(entry.end_time for entry in test_plan)
+
+            if max_finish_time <= self._objective:
+                # This works! Try with fewer resources
+                best_ncores = mid
+                best_plan = test_plan
+                best_graph = test_graph
+                upper_bound = mid - 1
+            else:
+                # Need more resources
+                lower_bound = mid + 1
+
+        if best_ncores is not None:
+            return best_ncores, best_plan, best_graph
+        return None
+
+    def _plan_with_qos_optimization(
+        self,
+        campaign: List[Workflow],
+        resource_requirements: Dict[int, Dict[str, float]],
+        requested_resources: int,
+    ) -> Tuple[List[PlanEntry], nx.DiGraph, str, int]:
+        """Find optimal QoS and resource allocation for the campaign.
+
+        Returns:
+            Tuple of (plan, graph, qos_name, ncores).
+        """
+        max_workflow_resources = self._get_max_ncores(resource_requirements)
+        upper_bound = requested_resources
+        lower_bound = max_workflow_resources
+
+        result = self._binary_search_resources(
+            campaign, resource_requirements, lower_bound, upper_bound
+        )
+
+        if result is not None:
+            ncores, plan, plan_graph = result
+            qos_candidate = self._find_suitable_qos_policies(requested_cores=ncores)
+            self._logger.info(f"Plan to execute {plan} with {ncores} cores")
+            return plan, plan_graph, qos_candidate.name, ncores
+        else:
+            raise ValueError(
+                f"Cannot meet {self._objective} min deadline with {requested_resources} cores. "
+                f"Please increase deadline or increase requested cores."
+            )
+
+    def _get_plan_graph(
+        self, plan: List[PlanEntry], resources: range
+    ) -> nx.DiGraph:
+        """Build dependency graph from the execution plan.
+
+        Args:
+            plan: Execution plan with scheduled workflows
+            resources: Available resource cores
+
+        Returns:
+            Directed acyclic graph representing workflow dependencies
+        """
+        self._logger.debug("Create resource dependency DAG")
+        graph = nx.DiGraph()
+
+        # Track which workflow is using each core
+        core_dependencies = {i: None for i in range(len(resources))}
+
+        for entry in plan:
+            # Find all previous tasks that occupied these cores
+            previous_tasks = {
+                core_dependencies[core]
+                for core in entry.cores
+                if core_dependencies[core] is not None
+            }
+
+            # Update core ownership
+            for core in entry.cores:
+                core_dependencies[core] = entry.workflow.id
+
+            # Add node and edges to graph
+            if not previous_tasks:
+                graph.add_node(entry.workflow.id)
+            else:
+                for predecessor_id in previous_tasks:
+                    graph.add_edge(predecessor_id, entry.workflow.id)
+
+        self._logger.info(f"Calculated graph {graph}")
+        return graph
+
+    def plan(
+        self,
+        campaign: List[Workflow] | None = None,
+        resource_requirements: Dict[int, Dict[str, float]] | None = None,
+        execution_schema: str | None = None,
+        requested_resources: int | None = None
+    ) -> Tuple[List[PlanEntry], nx.DiGraph, QosPolicy | None, int]:
+        """Plan campaign execution with resource allocation.
+
+        In batch mode, uses the requested resources directly.
+        In remote mode, performs QoS selection and binary search to find the minimum
+        resources that satisfy the campaign deadline.
+
+        Parameters
+        ----------
+        campaign : List[Workflow] | None
+            The campaign workflows to plan
+        resource_requirements : Dict[int, Dict[str, float]] | None
+            Resource requirements for each workflow
+        execution_schema : str | None
+            'batch' for fixed resources, 'remote' for optimized allocation
+        requested_resources : int | None
+            Number of cores
+
+        Returns
+        -------
+        Tuple[plan, graph, qos, ncores]
+            - plan: List of PlanEntry tuples
+            - graph: DAG representation of the campaign
+            - qos: QoS policy name (None for batch mode)
+            - ncores: Number of cores allocated
+        """
+        if execution_schema == "batch":
+            return self._plan_batch_mode(campaign, resource_requirements, requested_resources)
+        else:
+            return self._plan_with_qos_optimization(campaign, resource_requirements, requested_resources)
+
+    def _plan_batch_mode(
+        self,
+        campaign: List[Workflow],
+        resource_requirements: Dict[int, Dict[str, float]],
+        requested_resources: int
+    ) -> Tuple[List[PlanEntry], nx.DiGraph, None, int]:
+        """Plan execution for batch mode with fixed resources."""
+        plan, plan_graph = self._calculate_plan(
+            campaign=campaign,
+            resources=range(requested_resources),
+            resource_requirements=resource_requirements
+        )
+        self._logger.info(f"Plan to execute {plan} with {requested_resources} cores")
+        return plan, plan_graph, None, requested_resources
+
+    def _initialize_resource_estimates(
+        self, resource_requirements: Dict[int, Dict[str, float]]
+    ) -> None:
+        """Extract and store resource requirement estimates from workflows."""
+        self._estimated_walltime = []
+        self._estimated_cpus = []
+        self._estimated_memory = []
+
+        for resource_req in resource_requirements.values():
+            self._estimated_walltime.append(resource_req["req_walltime"])
+            self._estimated_cpus.append(resource_req["req_cpus"])
+            self._estimated_memory.append(resource_req["req_memory"])
+
+    def _get_sorted_workflow_indices(self) -> List[int]:
+        """Get workflow indices sorted by execution time (longest first).
+
+        Returns:
+            List of workflow indices in descending order of execution time
+        """
+        return [
+            idx for idx, _ in sorted(
+                enumerate(self._estimated_walltime),
+                key=lambda x: x[1],
+                reverse=True
+            )
+        ]
+
+    def _initialize_resource_free_times(
+        self, resources: range, start_time: float | int | list | np.ndarray
+    ) -> np.ndarray:
+        """Initialize array tracking when each resource becomes available.
+
+        Args:
+            resources: Range of available resource cores
+            start_time: Initial availability time(s)
+
+        Returns:
+            Array of availability times for each core
+        """
+        if isinstance(start_time, (np.ndarray, list)):
+            return np.array(start_time)
+        elif isinstance(start_time, (float, int)):
+            return np.array([start_time] * len(resources))
+        else:
+            return np.array([0.0] * len(resources))
+
+    def _find_best_resource_slot(
+        self,
+        workflow_idx: int,
+        resources: range,
+        resource_free: np.ndarray
+    ) -> Tuple[int, float]:
+        """Find the best resource slot for a workflow.
+
+        Args:
+            workflow_idx: Index of the workflow to schedule
+            resources: Available resource cores
+            resource_free: Array tracking when each core becomes available
+
+        Returns:
+            Tuple of (best_core_index, earliest_start_time)
+        """
+        walltime = self._estimated_walltime[workflow_idx]
+        memory_required = self._estimated_memory[workflow_idx]
+        cpus_required = self._estimated_cpus[workflow_idx]
+
+        min_end_time = float("inf")
+        best_core_idx = 0
+        core_idx = 0
+
+        while (core_idx + cpus_required) <= len(resources):
+            core_slice = slice(core_idx, core_idx + cpus_required)
+            start_time_candidate = resource_free[core_slice].max()
+            end_time_candidate = start_time_candidate + walltime
+
+            free_memory = self._get_free_memory(start_time_candidate, len(resources) / self._resources.cores_per_node)
+
+            if free_memory >= memory_required:
+                self._logger.debug(
+                    f"Workflow {workflow_idx}: candidate finish time {end_time_candidate}"
+                )
+                if end_time_candidate < min_end_time:
+                    min_end_time = end_time_candidate
+                    best_core_idx = core_idx
+                self._logger.debug(
+                    f"Workflow {workflow_idx}: minimum finish time {min_end_time}"
+                )
+            else:
+                self._logger.debug(
+                    f"Insufficient memory: {free_memory} MB available, {memory_required} MB required"
+                )
+
+            core_idx += cpus_required
+
+        return best_core_idx, min_end_time - walltime
+
+    def _calculate_plan(
+        self,
+        campaign: List[Workflow] | None = None,
+        resources: range | None = None,
+        resource_requirements: Dict[int, Dict[str, float]] | None = None,
+        start_time: float = 0.0,
+    ) -> Tuple[List[PlanEntry], nx.DiGraph]:
+        """Implement the core HEFT scheduling algorithm.
+
+        Args:
+            campaign: List of workflows to schedule
+            resources: Available resource cores
+            resource_requirements: Resource needs for each workflow
+            start_time: Initial time or per-core availability times
+
+        Returns:
+            Tuple of (execution_plan, dependency_graph)
+        """
+        # Use provided parameters or fall back to instance attributes
+        workflows = campaign if campaign else self._campaign
+        cores = (
+            resources
+            if resources
+            else range(self._resources.nodes * self._resources.cores_per_node)
+        )
+        requirements = resource_requirements if resource_requirements else self._resource_requirements
+
+        # Initialize estimation tables
+        self._initialize_resource_estimates(requirements)
+
+        # Reset plan for fresh scheduling
+        self._plan: List[PlanEntry] = []
+
+        # Sort workflows by execution time (longest first)
+        sorted_indices = self._get_sorted_workflow_indices()
+
+        # Track when each core becomes available
+        resource_free = self._initialize_resource_free_times(cores, start_time)
+
+        # Schedule each workflow
+        for workflow_idx in sorted_indices:
+            best_core_idx, start_time_actual = self._find_best_resource_slot(
+                workflow_idx, cores, resource_free
+            )
+
+            cpus_required = self._estimated_cpus[workflow_idx]
+            walltime = self._estimated_walltime[workflow_idx]
+            memory_required = self._estimated_memory[workflow_idx]
+            core_slice = slice(best_core_idx, best_core_idx + cpus_required)
+
+            # Create plan entry
+            plan_entry = PlanEntry(
+                workflow=workflows[workflow_idx],
+                cores=cores[core_slice],
+                memory=memory_required,
+                start_time=start_time_actual,
+                end_time=start_time_actual + walltime
+            )
+            self._plan.append(plan_entry)
+
+            # Update resource availability
+            resource_free[core_slice] = start_time_actual + walltime
+
+        # Build dependency graph
+        plan_graph = self._get_plan_graph(self._plan, cores)
+
+        # Sort plan by workflow ID for consistent ordering
+        self._plan = sorted(self._plan, key=lambda entry: entry.workflow.id)
+        self._logger.debug("Potential plan %s", self._plan)
+
+        return self._plan, plan_graph
+
+    def replan(
+        self,
+        campaign: List[Workflow] | None = None,
+        resources: range | None = None,
+        resource_requirements: Dict[int, Dict[str, float]] | None = None,
+        start_time: float = 0.0,
+    ) -> Tuple[List[PlanEntry], nx.DiGraph]:
+        """Recalculate the execution plan with updated parameters.
+
+        Args:
+            campaign: Updated list of workflows
+            resources: Updated resource allocation
+            resource_requirements: Updated resource requirements
+            start_time: New start time or per-core availability
+
+        Returns:
+            Tuple of (execution_plan, dependency_graph)
+        """
+        if campaign and resources and resource_requirements:
+            self._logger.debug("Replanning with updated parameters")
+            return self._calculate_plan(
+                campaign=campaign,
+                resources=resources,
+                resource_requirements=resource_requirements,
+                start_time=start_time,
+            )
+        else:
+            self._logger.debug("Nothing to replan - missing required parameters")
+            return self._plan, self._get_plan_graph(self._plan, resources or range(0))

socm/resources/__init__.py

@@ -0,0 +1,5 @@
+from .perlmutter import PerlmutterResource  # noqa: F401
+from .tiger import TigerResource  # noqa: F401
+from .universe import UniverseResource  # noqa: F401
+
+registered_resources = {"perlmutter": PerlmutterResource, "tiger3": TigerResource, "universe": UniverseResource}

socm/resources/perlmutter.py

@@ -0,0 +1,22 @@
+from ..core import QosPolicy, Resource
+
+
+class PerlmutterResource(Resource):
+    """
+    PerlmutterResource is a specialized Resource class for the Perlmutter HPC system.
+    It includes additional attributes specific to the Perlmutter system.
+    """
+
+    name: str = "perlmutter"
+    nodes: int = 3072
+    cores_per_node: int = 128
+    memory_per_node: int = 1000000  # in MB
+    default_qos: str = "regular"
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self.qos = [QosPolicy(name="regular", max_walltime=2880, max_jobs=5000, max_cores=393216),
+                    QosPolicy(name="interactive", max_walltime=240, max_jobs=2, max_cores=512),
+                    QosPolicy(name="shared_interactive", max_walltime=240, max_jobs=2, max_cores=64),
+                    QosPolicy(name="debug", max_walltime=30, max_jobs=5, max_cores=1024)
+                    ]

socm/resources/tiger.py

@@ -0,0 +1,24 @@
+from ..core import QosPolicy, Resource
+
+
+class TigerResource(Resource):
+    """
+    TigerResource is a specialized Resource class for the Tiger HPC system.
+    It includes additional attributes specific to the Tiger system.
+    """
+
+    name: str = "tiger3"
+    nodes: int = 492
+    cores_per_node: int = 112
+    memory_per_node: int = 1000000  # in MB
+    default_qos: str = "test"
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self.qos = [QosPolicy(name="test", max_walltime=60, max_jobs=1, max_cores=8000),
+                    QosPolicy(name="vshort", max_walltime=300, max_jobs=2000, max_cores=55104),
+                    QosPolicy(name="short", max_walltime=1440, max_jobs=50, max_cores=8000),
+                    QosPolicy(name="medium", max_walltime=4320, max_jobs=80, max_cores=4000),
+                    QosPolicy(name="long", max_walltime=8640, max_jobs=16, max_cores=1000),
+                    QosPolicy(name="vlong", max_walltime=21600, max_jobs=8, max_cores=900)
+                    ]

socm/resources/universe.py

@@ -0,0 +1,18 @@
+from ..core import QosPolicy, Resource
+
+
+class UniverseResource(Resource):
+    """
+    UniverseResource is a specialized Resource class for the Universe HPC system.
+    It includes additional attributes specific to the Universe system.
+    """
+
+    name: str = "universe"
+    nodes: int = 28
+    cores_per_node: int = 224
+    memory_per_node: int = 1000000  # in MB
+    default_qos: str = "main"
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self.qos = [QosPolicy(name="main", max_walltime=43200, max_jobs=5000, max_cores=6272)]

socm/utils/misc.py

@@ -0,0 +1,90 @@
+import ast
+from typing import Dict, List
+
+
+def parse_comma_separated_fields(config: dict, fields_to_parse: List[str]) -> dict:
+    """Convert comma-separated string values to lists."""
+    for key, value in config.items():
+        if isinstance(value, dict):
+            parse_comma_separated_fields(value, fields_to_parse)
+        elif key in fields_to_parse and isinstance(value, str) and ',' in value:
+            config[key] = [ast.literal_eval(item.strip()) for item in value.split(',')]
+    return config
+
+def get_workflow_entries(campaign_dict: dict, subcampaign_map: Dict[str, list] | None = None) -> Dict[str, dict]:
+    """
+    Extract workflow entries from a campaign dictionary using a predefined mapping.
+
+    Args:
+        campaign_dict: A dictionary containing campaign configuration
+        subcampaign_map: A dictionary mapping subcampaign names to lists of their workflow names
+                         E.g., {"ml-null-test": ["mission-tests", "wafer-tests"]}
+
+    Returns:
+        Dictionary containing workflow entries
+    """
+    campaign_data = campaign_dict.get("campaign", {})
+
+    # Default empty map if none provided
+    if subcampaign_map is None:
+        subcampaign_map = {}
+
+    # Collect all workflows (direct and from subcampaigns)
+    workflows = {}
+
+    for workflow_key, workflow_value in campaign_data.items():
+        # Skip non-dictionary values
+        if not isinstance(workflow_value, dict):
+            continue
+
+        # Check if this is a known subcampaign
+        if workflow_key in subcampaign_map:
+            # Process known workflows for this subcampaign
+            subcampaign_name = workflow_key
+            subcampaign_workflows = subcampaign_map[workflow_key]
+
+            # Create a copy of the subcampaign config without its workflows
+            subcampaign_common_config = {
+                k: v for k, v in workflow_value.items() if k not in subcampaign_workflows
+            }
+
+            for workflow_name in subcampaign_workflows:
+                if workflow_name in workflow_value:
+                    # Start with the workflow's own config
+                    workflow_config = workflow_value[workflow_name].copy()
+
+                    # Update with common subcampaign config
+                    workflow_config.update(subcampaign_common_config)
+
+                    if isinstance(workflow_config, dict):
+                        # Create combined key: subcampaign.workflow_name
+                        workflows[f"{subcampaign_name}.{workflow_name}"] = (
+                            workflow_config
+                        )
+        else:
+            # Treat as regular workflow
+            workflows[workflow_key] = workflow_value
+
+    return workflows
+
+
+def get_query_from_file(file_path: str) -> str:
+    """
+    Extract a SQL query from a file.
+
+    Args:
+        file_path: The path to the file containing the SQL query.
+
+    Returns:
+        The SQL query as a string.
+    """
+    query = "obs_id IN ("
+    with open(file_path, "r") as file:
+        obslist = file.readlines()
+        for obs_id in obslist:
+            obs_id = obs_id.strip()
+            query += f"'{obs_id}',"
+    query = query.rstrip(",")
+    query += ")"
+
+    return query

socm/utils/states.py

@@ -0,0 +1,17 @@
+# ------------------------------------------------------------------------------
+#  States
+from enum import Enum, auto
+
+
+class States(Enum):
+    """Workflow and campaign execution states."""
+
+    NEW = auto()  # New campaign is submitted
+    PLANNING = auto()  # Planning the execution of the campaign
+    EXECUTING = auto()  # At least one workflow is executing
+    DONE = auto()  # Campaign has finished successfully
+    FAILED = auto()  # Campaign execution has failed
+    CANCELED = auto()  # Campaign got canceled by the user.
+
+# Final states for a campaign
+CFINAL = [States.DONE, States.FAILED, States.CANCELED]

socm/workflows/__init__.py

@@ -0,0 +1,41 @@
+from socm.workflows.ml_mapmaking import MLMapmakingWorkflow  # noqa: F401
+from socm.workflows.ml_null_tests import (  # noqa: F401
+    DayNightNullTestWorkflow,
+    DirectionNullTestWorkflow,
+    ElevationNullTestWorkflow,
+    MoonCloseFarNullTestWorkflow,
+    MoonRiseSetNullTestWorkflow,
+    PWVNullTestWorkflow,
+    SunCloseFarNullTestWorkflow,
+    TimeNullTestWorkflow,
+    WaferNullTestWorkflow,
+)
+from socm.workflows.sat_simulation import SATSimWorkflow  # noqa: F401
+
+registered_workflows = {
+    "ml-mapmaking": MLMapmakingWorkflow,
+    "sat-sims": SATSimWorkflow,
+    "ml-null-tests.mission-tests": TimeNullTestWorkflow,
+    "ml-null-tests.wafer-tests": WaferNullTestWorkflow,
+    "ml-null-tests.direction-tests": DirectionNullTestWorkflow,
+    "ml-null-tests.pwv-tests": PWVNullTestWorkflow,
+    "ml-null-tests.day-night-tests": DayNightNullTestWorkflow,
+    "ml-null-tests.moonrise-set-tests": MoonRiseSetNullTestWorkflow,
+    "ml-null-tests.elevation-tests": ElevationNullTestWorkflow,
+    "ml-null-tests.sun-close-tests": SunCloseFarNullTestWorkflow,
+    "ml-null-tests.moon-close-tests": MoonCloseFarNullTestWorkflow,
+}
+
+subcampaign_map = {
+    "ml-null-tests": [
+        "mission-tests",
+        "wafer-tests",
+        "direction-tests",
+        "pwv-tests",
+        "day-night-tests",
+        "moonrise-set-tests",
+        "elevation-tests",
+        "sun-close-tests",
+        "moon-close-tests",
+    ]
+}