pyantz 0.1.0__py3-none-any.whl

pyantz/infrastructure/__init__.py

@@ -0,0 +1,4 @@
+"""
+Infrastructure provides the backbones of this runner. This module contains the code
+    to run all the various components to setup and execute the jobs in the configuration
+"""

pyantz/infrastructure/config/base.py

@@ -0,0 +1,332 @@
+"""
+This is the base level of the configuration for the core components
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from collections.abc import Mapping
+from typing import Any, Callable, Literal, TypeAlias, Union
+
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    ConfigDict,
+    Field,
+    JsonValue,
+    field_serializer,
+    model_validator,
+    validate_call,
+)
+from typing_extensions import Annotated
+
+from pyantz.infrastructure.config.get_functions import (
+    get_function_by_name_strongly_typed,
+    get_params_model,
+    set_job_type,
+    set_params_model,
+)
+from pyantz.infrastructure.core.status import Status
+from pyantz.infrastructure.core.variables import is_variable
+
+from .local_submitter import LocalSubmitterConfig
+from .submitters.slurm_submitter import SlurmBasicSubmitter
+
+PrimitiveType: TypeAlias = str | int | float | bool | None
+AntzConfig: TypeAlias = Union[
+    "Config", "PipelineConfig", "JobConfig", "SubmitterJobConfig", "MutableJobConfig"
+]
+ParametersType: TypeAlias = (
+    Mapping[str, AntzConfig | list[AntzConfig] | JsonValue] | None
+)
+SubmitFunctionType: TypeAlias = Callable[["Config"], None]
+JobFunctionType: TypeAlias = Callable[
+    ["ParametersType", logging.Logger],
+    Status,
+]
+SubmitterJobFunctionType: TypeAlias = Callable[
+    [
+        "ParametersType",
+        SubmitFunctionType,
+        Mapping[str, PrimitiveType],
+        "PipelineConfig",
+        logging.Logger,
+    ],
+    Status,
+]
+MutableJobFunctionType: TypeAlias = Callable[
+    ["ParametersType", Mapping[str, PrimitiveType], logging.Logger],
+    tuple[
+        Status,
+        Mapping[str, PrimitiveType],
+    ],
+]
+
+
+class _AbstractJobConfig(BaseModel, frozen=True):
+    """holds common functions for the various job configs"""
+
+    function: Callable[..., Any]
+    parameters: ParametersType
+    name: str = "some job"
+    id: uuid.UUID = Field(default_factory=uuid.uuid4, validate_default=True)
+
+    @field_serializer("function")
+    def serialize_function(self, func: JobFunctionType, info):
+        """To serialize function, store the import path to the func
+        instead of its handle as a str
+        """
+        if hasattr(func, "__wrapped__"):
+            return self.serialize_function(func.__wrapped__, info)
+        return func.__module__ + "." + func.__name__
+
+    @model_validator(mode="after")
+    def check_parameters_match(self: "_AbstractJobConfig") -> "_AbstractJobConfig":
+        """Checks that the config parameters match the expected parameters for the function"""
+
+        params_model = get_params_model(self.function)
+        if params_model is None:
+            return self
+
+        if params_model is None:
+            return self
+        if not isinstance(params_model, type) or not issubclass(
+            params_model, BaseModel
+        ):
+            raise ValueError(
+                f"Invalid parameters mode for function {self.function.__name__}"
+            )
+
+        # If the parameters are None or not a mapping, error in validation
+        if self.parameters is None:
+            raise ValueError(
+                f"Parameters cannot be None for function {self.function.__name__}"
+            )
+        if not isinstance(self.parameters, Mapping):
+            raise ValueError(
+                f"Parameters must be a mapping for function {self.function.__name__}"
+            )
+
+        if any(isinstance(field, BaseModel) for field in self.parameters.values()):
+            # cannot check jobs or pipelines
+            return self
+        if any(
+            is_variable(field)
+            for field in self.parameters.values()
+            if isinstance(field, PrimitiveType)
+        ):
+            return self
+        if not params_model.model_validate(self.parameters):
+            raise ValueError(
+                f"Parameters do not match expected parameters for function {self.function.__name__}"
+            )
+        return self
+
+
+class MutableJobConfig(_AbstractJobConfig, frozen=True):
+    """Configuration of a submitter job, with different function param types
+    These jobs gain access to the submit function and can submit
+        entirely new pipelines of execution
+
+    However, they must ALWAYS BE FINAL
+    """
+
+    type: Literal["mutable_job"]
+    function: Annotated[
+        MutableJobFunctionType,
+        BeforeValidator(get_function_by_name_strongly_typed("mutable")),
+    ]
+
+
+class SubmitterJobConfig(_AbstractJobConfig, frozen=True):
+    """Configuration of a submitter job, with different function param types
+    These jobs gain access to the submit function and can submit
+        entirely new pipelines of execution
+
+    However, they must ALWAYS BE FINAL
+    """
+
+    type: Literal["submitter_job"]
+    function: Annotated[
+        SubmitterJobFunctionType,
+        BeforeValidator(get_function_by_name_strongly_typed("submitter")),
+    ]
+
+
+class JobConfig(_AbstractJobConfig, frozen=True):
+    """Configuration of a job"""
+
+    type: Literal["job"] | Literal["simple_job"]
+    function: Annotated[
+        JobFunctionType,
+        BeforeValidator(get_function_by_name_strongly_typed("simple")),
+    ]
+
+
+class PipelineConfig(BaseModel, frozen=True):
+    """Configuration of a pipeline, which is a series of jobs or sub-pipelines"""
+
+    type: Literal["pipeline"]
+    name: str = "pipeline"
+    curr_stage: int = 0
+    id: uuid.UUID = Field(default_factory=uuid.uuid4, validate_default=True)
+    status: int = Status.READY
+    max_allowed_restarts: int = 0
+    curr_restarts: int = 0
+    stages: list[
+        Annotated[
+            Union[MutableJobConfig, SubmitterJobConfig, JobConfig],
+            Field(discriminator="type"),
+        ]
+    ]
+    # stages: list[JobConfig | SubmitterJobConfig | MutableJobConfig]
+
+
+class LoggingConfig(BaseModel, frozen=True):
+    """The configuration of logging"""
+
+    type: Literal["off", "file", "console", "remote"] = (
+        "console"  # default to logging to screen
+    )
+    level: int = logging.CRITICAL  # default to only logging on crashes
+    directory: str | None = "./log"
+
+
+class Config(BaseModel, frozen=True):
+    """The global configuration to submit to runner"""
+
+    variables: Mapping[str, PrimitiveType]
+    config: PipelineConfig
+
+
+class InitialConfig(BaseModel, frozen=True):
+    """The configuration of both the jobs and the submitters"""
+
+    analysis_config: Config
+    submitter_config: LocalSubmitterConfig | SlurmBasicSubmitter = Field(
+        discriminator="type"
+    )
+    logging_config: LoggingConfig = LoggingConfig()
+
+
+def mutable_job(
+    params_model: type[BaseModel] | None,
+) -> Callable[
+    [
+        Callable[
+            ["ParametersType", Mapping[str, PrimitiveType], logging.Logger],
+            tuple[Status, Mapping[str, PrimitiveType]],
+        ]
+    ],
+    MutableJobFunctionType,
+]:
+    """Wrap a mutable job and add its parameters model to its definition to allow
+    the configuration parser to validate the parameters
+    """
+
+    def mark_mutable_job(
+        fn: Callable[
+            ["ParametersType", Mapping[str, PrimitiveType], logging.Logger],
+            tuple[Status, Mapping[str, PrimitiveType]],
+        ],
+    ) -> MutableJobFunctionType:
+        """Wrap a mutable job to
+        1. Allow it to accept variable args if a user incorrectly marks job
+        2. Allow for type checking in the pydantic model
+        """
+
+        @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
+        def _mutable_job(
+            params: ParametersType,
+            variables: Mapping[str, PrimitiveType],
+            logger: logging.Logger,
+        ) -> tuple[Status, Mapping[str, PrimitiveType]]:
+            return fn(params, variables, logger)
+
+        set_job_type(_mutable_job, "mutable")
+        set_params_model(_mutable_job, params_model)
+        setattr(_mutable_job, "__wrapped__", fn)
+        return _mutable_job
+
+    return mark_mutable_job
+
+
+def submitter_job(
+    params_model: type[BaseModel] | None,
+) -> Callable[[SubmitterJobFunctionType], SubmitterJobFunctionType]:
+    """Wrap a submitter job and add its parameters model to its definition to allow
+    the configuration parser to validate the parameters
+    """
+
+    def mark_submitter_job(fn: SubmitterJobFunctionType) -> SubmitterJobFunctionType:
+        """Wrap a submitter job to
+        1. Allow it to accept variable args if a user incorrectly marks job
+        2. Allow for type checking in the pydantic model
+        """
+
+        @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
+        def _submitter_job(
+            params: ParametersType,
+            submitter: SubmitFunctionType,
+            variables: Mapping[str, PrimitiveType],
+            pipeline_config: PipelineConfig,
+            logger: logging.Logger,
+        ) -> Status:
+            return fn(params, submitter, variables, pipeline_config, logger)
+
+        set_job_type(_submitter_job, "submitter")
+        set_params_model(_submitter_job, params_model)
+        setattr(_submitter_job, "__wrapped__", fn)
+        return _submitter_job
+
+    return mark_submitter_job
+
+
+def simple_job(
+    params_model: type[BaseModel] | None,
+) -> Callable[[Callable[[ParametersType, logging.Logger], Status]], JobFunctionType]:
+    """Wrap a simple job and add its parameters model to its definition to allow
+    the configuration parser to validate the parameters
+    """
+
+    def mark_simple_job(
+        fn: Callable[[ParametersType, logging.Logger], Status],
+    ) -> JobFunctionType:
+        """Wrap a simple job to
+        1. Allow it to accept variable args if a user incorrectly marks job
+        2. Allow for type checking in the pydantic model
+        """
+
+        @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
+        def _simple_job(
+            params: ParametersType,
+            logger: logging.Logger,
+        ) -> Status:
+            return fn(params, logger)
+
+        set_job_type(_simple_job, "simple")
+        set_params_model(_simple_job, params_model)
+        setattr(_simple_job, "__wrapped__", fn)
+        return _simple_job
+
+    return mark_simple_job
+
+
+__all__ = [
+    "simple_job",
+    "submitter_job",
+    "mutable_job",
+    "InitialConfig",
+    "Config",
+    "LoggingConfig",
+    "PipelineConfig",
+    "JobConfig",
+    "SubmitterJobConfig",
+    "MutableJobConfig",
+    "ParametersType",
+    "PrimitiveType",
+    "SubmitFunctionType",
+    "MutableJobFunctionType",
+    "JobFunctionType",
+]

pyantz/infrastructure/config/get_functions.py

@@ -0,0 +1,132 @@
+"""Functions to dynamically import and tag functions from a configuration"""
+
+import importlib
+from typing import Any, Callable
+
+from pydantic import BaseModel
+
+_PYANTZ_JOB_TYPE_FIELD: str = "__pyantz_job_type__"
+_PYANTZ_PARAMS_MODEL_FIELD: str = "__pyantz_param_model__"
+
+
+def set_params_model(
+    fn: Callable[..., Any], params_model: type[BaseModel] | None
+) -> Callable[..., Any]:
+    """Sets the parameters model field"""
+    setattr(fn, _PYANTZ_PARAMS_MODEL_FIELD, params_model)
+    return fn
+
+
+def set_job_type(fn: Callable[..., Any], job_type: str) -> Callable[..., Any]:
+    """Set the job type field"""
+    setattr(fn, _PYANTZ_JOB_TYPE_FIELD, job_type)
+    return fn
+
+
+def get_params_model(fn: Callable[..., Any]) -> type[BaseModel] | None:
+    """For the provided callable, if it has a params model field,
+    return that field for type checking on instatiation of the job
+    """
+    if hasattr(fn, _PYANTZ_PARAMS_MODEL_FIELD):
+        return getattr(fn, _PYANTZ_PARAMS_MODEL_FIELD)
+    return None
+
+
+def get_job_type(fn: Callable[..., Any] | None) -> str | None:
+    """For a provided callable, return what type of job it is
+
+    This API is guaranteed to be stable; our implementation of how
+        to mark functions is not. SO **USE THIS** to check
+
+    :param fn: any function which may or may not be marked
+    :type fn: Callable[..., Any]
+    :return: if the function is marked, return the mark type; else None
+    :rtype: str | None
+    """
+    if fn is None:
+        return fn
+    if hasattr(fn, _PYANTZ_JOB_TYPE_FIELD):
+        return getattr(fn, _PYANTZ_JOB_TYPE_FIELD)
+    return None
+
+
+def get_function_by_name_strongly_typed(
+    func_type_name: str | tuple[str, ...], strict: bool | None = None
+) -> Callable[[Any], Callable[..., Any] | None]:
+    """Returns a function Calls get_function_by_name and checks that the function type is correct
+
+    Uses strict rules for internal functions; otherwise uses non-strict
+        can be overriden with the strict argument
+    If strict is True,
+        requires that the function is wrapped in the correct wrapper from job_decorators.py
+    if strict is false,
+        if the function is not wrapped in any of those wrappers, will skip checking
+
+    Args:
+        func_type_name: the name of the wrapper in job_decorators
+        strict: overrides the default behavior if provided, see notes above
+    """
+    # strict for PyAntz jobs because we should at least be consistent!
+    if strict is None:
+        if isinstance(func_type_name, str):
+            strict = func_type_name.startswith("pyantz")
+        else:
+            strict = all(name.startswith("pyantz") for name in func_type_name)
+
+    def typed_get_function_by_name(
+        func_name_or_any: Any,
+    ) -> Callable[..., Any] | None:
+        func_handle = get_function_by_name(func_name_or_any)
+        job_type = get_job_type(func_handle)
+        if job_type is None and strict:
+            return None
+        if job_type is None:
+            return func_handle
+        if isinstance(func_type_name, str):
+            if job_type != func_type_name:
+                return None
+            return func_handle
+        if job_type in func_type_name:
+            return func_handle
+        return None
+
+    return typed_get_function_by_name
+
+
+def get_function_by_name(func_name_or_any: Any) -> Callable[..., Any] | None:
+    """Links to the function described by config
+
+    Args:
+        config (JobConfig): configuration of the job to link
+
+    Returns:
+        Callable[[ParametersType, Callable[[PipelineConfig], None]], Status] } None:
+            a function that takes parameters and a
+            submitter callable and returns a status after executing
+            Returns None if it is unable to find the correct function
+
+    """
+
+    if not isinstance(func_name_or_any, str):
+        return None
+
+    name: str = func_name_or_any
+
+    components = name.split(".")
+    func_name = components[-1]
+    mod_name = ".".join(components[:-1])
+
+    try:
+        mod = importlib.import_module(mod_name)
+    except ModuleNotFoundError as _:
+        return None
+
+    if not hasattr(mod, func_name):
+        return None
+
+    func = getattr(mod, func_name)
+
+    if not callable(func):
+        return None
+
+    return func

pyantz/infrastructure/config/local_submitter.py

@@ -0,0 +1,20 @@
+"""Configuration for the Local Submitter
+
+num_concurrent_jobs controls how many processes to spawn for the manager
+"""
+
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class LocalSubmitterConfig(BaseModel, frozen=True):
+    """
+    The configuration of the local submitter
+
+    num_concurrent_jobs (int): number of processes to run jobs
+    """
+
+    type: Literal["local"]
+    name: str = "local submitter"
+    num_concurrent_jobs: int = 1

pyantz/infrastructure/config/submitters/slurm_submitter.py

@@ -0,0 +1,50 @@
+"""Configuration of the slurm submitter"""
+
+from typing import Literal
+import enum
+
+from pydantic import BaseModel, DirectoryPath, Field
+
+
+class RetryPolicy(enum.StrEnum):
+    """Retry policies are how a slurm submitter should try to resubmit jobs that failed
+
+    naive: just resubmit with no changes
+    exclude: exclude the node that failed previously
+    include: only submit to nodes with previously successful jobs
+    current: only submit to the current node
+    """
+
+    NAIVE = "naive"
+    EXCLUDE = "exclude"
+    INCLUDE = "include"
+    CURRENT = "current"
+
+
+class SlurmBasicSubmitter(BaseModel, frozen=True):
+    """Configuration of the basic slurm submitter (basic_slurm)
+
+    Fields:
+    - type (str): always set to slurm basic
+    - name (str): whatever name you want for the submitters
+    - max_submit_retries (int): when submitting, how many times to retry a failed submission
+        useful if you think the failure is from random node issues
+        defaults to 0
+    - retry policy: see RetryPolicy enum docstring
+        how a submitter should resubmit failed jobs
+        defaults to NAIVE
+    - submit_wait_time (int): how long after submitting to wait for job fail/success
+        in seconds
+        defaults to 3
+    - slurm_command (str): the command to use to submit. for now, must be sbatch
+        in the future, could come up with an srun solution
+    """
+
+    type: Literal["slurm_basic"]
+    name: str = "basic slurm submitter"
+    max_submit_retries: int = 0
+    retry_policy: RetryPolicy = RetryPolicy.NAIVE
+    submit_wait_time: int = 3
+    slurm_command: Literal["sbatch"] = "sbatch"
+    working_directory: DirectoryPath
+    grid_cmd_args: list[str] = Field(..., default_factory=lambda: [])  # type: ignore

pyantz/infrastructure/core/job.py

@@ -0,0 +1,42 @@
+"""A job is the basic unit of execution in this module
+
+Each job performs one user-assigned task and returns its state.
+"""
+
+# pylint: disable=duplicate-code
+
+import logging
+from collections.abc import Mapping
+
+from pyantz.infrastructure.config.base import JobConfig, PrimitiveType
+from pyantz.infrastructure.core.status import Status
+from pyantz.infrastructure.core.variables import resolve_variables
+
+
+def run_job(
+    config: JobConfig,
+    variables: Mapping[str, PrimitiveType],
+    logger: logging.Logger,
+) -> Status:
+    """Run a job, which is the smallest atomic task of antz"""
+    status: Status
+    func_handle = config.function
+    logger.debug("Running job %s, with func handle: %s", config.id, str(func_handle))
+
+    params = resolve_variables(config.parameters, variables)
+    logger.debug("Running function with parameters %s", str(params))
+
+    try:
+        ret = func_handle(params, logger)
+        if isinstance(ret, Status):
+            status = ret
+        else:
+            logger.warning(
+                "Return of function was not an ANTZ status, this is an automatic error"
+            )
+            status = Status.ERROR  # bad return type is an error
+    except Exception as exc:  # pylint: disable=broad-exception-caught
+        logger.warning("Unexpected error", exc_info=exc)
+        status = Status.ERROR
+    logger.debug("Finished job %s with status %s", config.id, str(status))
+    return status

pyantz/infrastructure/core/manager.py

@@ -0,0 +1,20 @@
+"""Run a general config, which is a pipeline with a scope (variables)"""
+
+import logging
+from typing import Callable
+
+from pyantz.infrastructure.config.base import Config
+from pyantz.infrastructure.core.pipeline import run_pipeline
+
+
+def run_manager(
+    config: Config, submit_fn: Callable[[Config], None], logger: logging.Logger
+) -> None:
+    """Run the configuration"""
+    logger.debug("Manager starting up pipeline with id %d", config.config.id)
+    run_pipeline(
+        config=config.config,
+        variables=config.variables,
+        submit_fn=submit_fn,
+        logger=logger,
+    )

pyantz/infrastructure/core/mutable_job.py

@@ -0,0 +1,44 @@
+"""Mutable jobs allow the function to edit the variables of the outer scope"""
+
+# pylint: disable=duplicate-code
+
+import logging
+from collections.abc import Mapping
+from copy import deepcopy
+
+from pyantz.infrastructure.config.base import MutableJobConfig, PrimitiveType
+from pyantz.infrastructure.core.status import Status
+from pyantz.infrastructure.core.variables import resolve_variables
+
+
+def run_mutable_job(
+    config: MutableJobConfig,
+    variables: Mapping[str, PrimitiveType],
+    logger: logging.Logger,
+) -> tuple[Status, Mapping[str, PrimitiveType]]:
+    """Run a job, which is the smallest atomic task of antz"""
+
+    status: Status
+    func_handle = config.function
+    logger.debug("Running job %s, with func handle: %s", config.id, str(func_handle))
+
+    params = resolve_variables(config.parameters, variables)
+    logger.debug("Running function with parameters %s", str(params))
+
+    try:
+        ret_status, ret_vars = func_handle(params, deepcopy(variables), logger)
+        if isinstance(ret_status, Status):
+            status = ret_status
+        else:
+            logger.warning(
+                "Return of function was not an ANTZ status, this is an automatic error"
+            )
+            status = Status.ERROR  # bad return type is an error
+    except Exception as exc:  # pylint: disable=broad-exception-caught
+        logger.warning("Unexpected error", exc_info=exc)
+        status = Status.ERROR
+    logger.debug("Finished job %s with status %s", config.id, str(status))
+
+    if status == Status.ERROR:
+        return status, variables
+    return status, ret_vars