Merge pull request #1552 from Libensemble/refactor/310_plus_typehints

Refactor/310 plus typehints

Author: jlnav

docs/data_structures/libE_specs.rst

@@ -205,14 +205,14 @@ libEnsemble is primarily customized by setting options within a ``LibeSpecs`` cl
                 **save_H_and_persis_on_abort** [bool] = ``True``:
                     Save states of ``H`` and ``persis_info`` to file on aborting after an exception.
 
-                **save_H_on_completion** Optional[bool] = ``False``
+                **save_H_on_completion** bool | None = ``False``
                     Save state of ``H`` to file upon completing a workflow. Also enabled when either ``save_every_k_sims``
                     or ``save_every_k_gens`` is set.
 
-                **save_H_with_date** Optional[bool] = ``False``
+                **save_H_with_date** bool | None = ``False``
                     Save ``H`` filename contains date and timestamp.
 
-                **H_file_prefix** Optional[str] = ``"libE_history"``
+                **H_file_prefix** str | None = ``"libE_history"``
                     Prefix for ``H`` filename.
 
                 **use_persis_return_gen** [bool] = ``False``:

libensemble/alloc_funcs/give_sim_work_first.py

@@ -1,5 +1,4 @@
 import time
-from typing import Tuple
 
 import numpy as np
 import numpy.typing as npt
@@ -15,7 +14,7 @@ def give_sim_work_first(
     alloc_specs: dict,
     persis_info: dict,
     libE_info: dict,
-) -> Tuple[dict]:
+) -> tuple[dict]:
     """
     Decide what should be given to workers. This allocation function gives any
     available simulation work first, and only when all simulations are

libensemble/ensemble.py

@@ -1,7 +1,6 @@
 import importlib
 import json
 import logging
-from typing import Optional
 
 import numpy.typing as npt
 import tomli
@@ -270,15 +269,15 @@ class Ensemble:
 
     def __init__(
         self,
-        sim_specs: Optional[SimSpecs] = SimSpecs(),
-        gen_specs: Optional[GenSpecs] = GenSpecs(),
-        exit_criteria: Optional[ExitCriteria] = {},
-        libE_specs: Optional[LibeSpecs] = LibeSpecs(),
-        alloc_specs: Optional[AllocSpecs] = AllocSpecs(),
-        persis_info: Optional[dict] = {},
-        executor: Optional[Executor] = None,
-        H0: Optional[npt.NDArray] = None,
-        parse_args: Optional[bool] = False,
+        sim_specs: SimSpecs | None = SimSpecs(),
+        gen_specs: GenSpecs | None = GenSpecs(),
+        exit_criteria: ExitCriteria | None = {},
+        libE_specs: LibeSpecs | None = LibeSpecs(),
+        alloc_specs: AllocSpecs | None = AllocSpecs(),
+        persis_info: dict | None = {},
+        executor: Executor | None = None,
+        H0: npt.NDArray | None = None,
+        parse_args: bool | None = False,
     ):
         self.sim_specs = sim_specs
         self.gen_specs = gen_specs

libensemble/executors/balsam_executor.py

@@ -74,11 +74,12 @@ class HelloApp(ApplicationDefinition):
 .. _Globus: https://www.globus.org/
 """
 
+from __future__ import annotations
+
 import datetime
 import logging
 import os
 import time
-from typing import Any, Dict, List, Optional, Union
 
 from balsam import util
 
@@ -106,9 +107,9 @@ class BalsamTask(Task):
 
     def __init__(
         self,
-        app: Optional[Application] = None,
+        app: Application | None = None,
         app_args: dict = None,
-        workdir: Optional[str] = None,
+        workdir: str | None = None,
         stdout: str = None,
         stderr: str = None,
         workerid: int = None,
@@ -122,7 +123,7 @@ def __init__(
         # May want to override workdir with Balsam value when it exists
         Task.__init__(self, app, app_args, workdir, stdout, stderr, workerid)
 
-    def _get_time_since_balsam_submit(self) -> Union[int, float]:
+    def _get_time_since_balsam_submit(self) -> int:
         """Return time since balsam task entered ``RUNNING`` state"""
         event_query = EventLog.objects.filter(job_id=self.process.id, to_state="RUNNING")
         if not len(event_query):
@@ -203,7 +204,7 @@ def poll(self) -> None:
             self.state = "FAILED"
             self._set_complete()
 
-    def wait(self, timeout: Optional[int] = None) -> None:
+    def wait(self, timeout: int | None = None) -> None:
         """Waits on completion of the task or raises ``TimeoutExpired``.
 
         Status attributes of task are updated on completion.
@@ -280,10 +281,10 @@ def add_app(self, *args) -> None:
     def register_app(
         self,
         BalsamApp: ApplicationDefinition,
-        app_name: Optional[str] = None,
-        calc_type: Optional[str] = None,
+        app_name: str | None = None,
+        calc_type: str | None = None,
         desc: str = None,
-        precedent: Optional[str] = None,
+        precedent: str | None = None,
     ) -> None:
         """Registers a Balsam ``ApplicationDefinition`` to libEnsemble. This class
         instance *must* have a ``site`` and ``command_template`` specified. See
@@ -331,9 +332,9 @@ def submit_allocation(
         job_mode: str = "mpi",
         queue: str = "local",
         project: str = "local",
-        optional_params: Dict[Any, Any] = {},
-        filter_tags: Dict[Any, Any] = {},
-        partitions: List[Any] = [],
+        optional_params: dict = {},
+        filter_tags: dict = {},
+        partitions: list = [],
     ) -> BatchJob:
         """
         Submits a Balsam ``BatchJob`` machine allocation request to Balsam.
@@ -435,14 +436,14 @@ def set_resources(self, resources: str) -> None:
 
     def submit(
         self,
-        calc_type: Optional[str] = None,
-        app_name: Optional[str] = None,
+        calc_type: str | None = None,
+        app_name: str | None = None,
         app_args: dict = None,
         num_procs: int = None,
         num_nodes: int = None,
         procs_per_node: int = None,
         max_tasks_per_node: int = None,
-        machinefile: Optional[str] = None,
+        machinefile: str | None = None,
         gpus_per_rank: int = 0,
         transfers: dict = {},
         workdir: str = "",

libensemble/executors/executor.py

@@ -12,7 +12,7 @@
 import sys
 import time
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any
 
 import libensemble.utils.launcher as launcher
 from libensemble.message_numbers import (
@@ -24,7 +24,6 @@
     WORKER_DONE,
     WORKER_KILL_ON_TIMEOUT,
 )
-from libensemble.resources.resources import Resources
 from libensemble.utils.timer import TaskTimer
 
 logger = logging.getLogger(__name__)
@@ -78,10 +77,10 @@ class Application:
     def __init__(
         self,
         full_path: str,
-        name: Optional[str] = None,
-        calc_type: Optional[str] = "sim",
-        desc: Optional[str] = None,
-        pyobj: Optional[Any] = None,  # used by balsam_executor to store ApplicationDefinition
+        name: str | None = None,
+        calc_type: str | None = "sim",
+        desc: str | None = None,
+        pyobj: Any | None = None,  # used by balsam_executor to store ApplicationDefinition
         precedent: str = "",
     ) -> None:
         """Instantiates a new Application instance."""
@@ -101,7 +100,7 @@ def __init__(
         self.app_cmd = " ".join(filter(None, [self.precedent, self.full_path]))
 
 
-def jassert(test: Optional[Union[Application, bool]], *args) -> None:
+def jassert(test: Application | bool | None, *args) -> None:
     "Version of assert that raises a ExecutorException"
     if not test:
         raise ExecutorException(*args)
@@ -170,7 +169,7 @@ def _add_to_env(self, key, value):
         """Add to task environment - overwrites if already set"""
         self.env[key] = value
 
-    def workdir_exists(self) -> Optional[bool]:
+    def workdir_exists(self) -> bool | None:
         """Returns true if the task's workdir exists"""
         return self.workdir and os.path.exists(self.workdir)
 
@@ -260,7 +259,7 @@ def poll(self) -> None:
 
         self._set_complete()
 
-    def wait(self, timeout: Optional[float] = None) -> None:
+    def wait(self, timeout: float | None = None) -> None:
         """Waits on completion of the task or raises TimeoutExpired exception
 
         Status attributes of task are updated on completion.
@@ -288,7 +287,7 @@ def wait(self, timeout: Optional[float] = None) -> None:
 
         self._set_complete()
 
-    def result(self, timeout: Optional[Union[int, float]] = None) -> str:
+    def result(self, timeout: int | float | None = None) -> str:
         """Wrapper for task.wait() that also returns the task's status on completion.
 
         Parameters
@@ -303,7 +302,7 @@ def result(self, timeout: Optional[Union[int, float]] = None) -> str:
         self.wait(timeout=timeout)
         return self.state
 
-    def exception(self, timeout: Optional[Union[int, float]] = None):
+    def exception(self, timeout: int | float | None = None):
         """Wrapper for task.wait() that instead returns the task's error code on completion.
 
         Parameters
@@ -386,7 +385,7 @@ class Executor:
 
     executor = None
 
-    def _wait_on_start(self, task: Task, fail_time: Optional[int] = None) -> None:
+    def _wait_on_start(self, task: Task, fail_time: int | None = None) -> None:
         """Called by submit when wait_on_start is True.
 
         Blocks until task polls as having started.
@@ -472,7 +471,7 @@ def default_app(self, calc_type: str) -> Application:
         jassert(app, f"Default {calc_type} app is not set")
         return app
 
-    def set_resources(self, resources: Resources):
+    def set_resources(self, resources):
         # Does not use resources
         pass
 
@@ -493,9 +492,9 @@ def set_gen_procs_gpus(self, libE_info):
     def register_app(
         self,
         full_path: str,
-        app_name: Optional[str] = None,
-        calc_type: Optional[str] = None,
-        desc: Optional[str] = None,
+        app_name: str | None = None,
+        calc_type: str | None = None,
+        desc: str | None = None,
         precedent: str = "",
     ) -> None:
         """Registers a user application to libEnsemble.
@@ -571,7 +570,7 @@ def manager_kill_received(self) -> bool:
         return False
 
     def polling_loop(
-        self, task: Task, timeout: Optional[int] = None, delay: float = 0.1, poll_manager: bool = False
+        self, task: Task, timeout: int | None = None, delay: float = 0.1, poll_manager: bool = False
     ) -> int:
         """Optional, blocking, generic task status polling loop. Operates until the task
         finishes, times out, or is optionally killed via a manager signal. On completion, returns a
@@ -637,7 +636,7 @@ def polling_loop(
 
         return calc_status
 
-    def get_task(self, taskid: Union[str, int]) -> Optional[Task]:
+    def get_task(self, taskid: str | int) -> Task | None:
         """Returns the task object for the supplied task ID"""
         task = next((j for j in self.list_of_tasks if j.id == taskid), None)
         if task is None:
@@ -681,14 +680,14 @@ def _check_app_exists(self, full_path: str) -> None:
 
     def submit(
         self,
-        calc_type: Optional[str] = None,
-        app_name: Optional[str] = None,
-        app_args: Optional[str] = None,
-        stdout: Optional[str] = None,
-        stderr: Optional[str] = None,
-        dry_run: Optional[bool] = False,
-        wait_on_start: Optional[bool] = False,
-        env_script: Optional[str] = None,
+        calc_type: str | None = None,
+        app_name: str | None = None,
+        app_args: str | None = None,
+        stdout: str | None = None,
+        stderr: str | None = None,
+        dry_run: bool | None = False,
+        wait_on_start: bool | None = False,
+        env_script: str | None = None,
     ) -> Task:
         """Create a new task and run as a local serial subprocess.
 

libensemble/executors/mpi_executor.py

@@ -15,13 +15,11 @@
 import logging
 import os
 import time
-from typing import List, Optional, Union
 
 import libensemble.utils.launcher as launcher
 from libensemble.executors.executor import Executor, ExecutorException, Task
 from libensemble.executors.mpi_runner import MPIRunner
 from libensemble.resources.mpi_resources import get_MPI_variant
-from libensemble.resources.resources import Resources
 
 logger = logging.getLogger(__name__)
 # To change logging level for just this module
@@ -137,11 +135,11 @@ def set_gen_procs_gpus(self, libE_info):
         self.gen_nprocs = libE_info.get("num_procs")
         self.gen_ngpus = libE_info.get("num_gpus")
 
-    def set_resources(self, resources: Resources) -> None:
+    def set_resources(self, resources) -> None:
         self.resources = resources
 
     def _launch_with_retries(
-        self, task: Task, subgroup_launch: bool, wait_on_start: bool, run_cmd: List[str], use_shell: bool
+        self, task: Task, subgroup_launch: bool, wait_on_start: bool, run_cmd: list[str], use_shell: bool
     ) -> None:
         """Launch task with retry mechanism"""
         retry_count = 0
@@ -189,25 +187,25 @@ def _launch_with_retries(
 
     def submit(
         self,
-        calc_type: Optional[str] = None,
-        app_name: Optional[str] = None,
-        num_procs: Optional[int] = None,
-        num_nodes: Optional[int] = None,
-        procs_per_node: Optional[int] = None,
-        num_gpus: Optional[int] = None,
-        machinefile: Optional[str] = None,
-        app_args: Optional[str] = None,
-        stdout: Optional[str] = None,
-        stderr: Optional[str] = None,
-        stage_inout: Optional[str] = None,
-        hyperthreads: Optional[bool] = False,
-        dry_run: Optional[bool] = False,
-        wait_on_start: Optional[bool] = False,
-        extra_args: Optional[str] = None,
-        auto_assign_gpus: Optional[bool] = False,
-        match_procs_to_gpus: Optional[bool] = False,
-        env_script: Optional[str] = None,
-        mpi_runner_type: Optional[Union[str, dict]] = None,
+        calc_type: str | None = None,
+        app_name: str | None = None,
+        num_procs: int | None = None,
+        num_nodes: int | None = None,
+        procs_per_node: int | None = None,
+        num_gpus: int | None = None,
+        machinefile: str | None = None,
+        app_args: str | None = None,
+        stdout: str | None = None,
+        stderr: str | None = None,
+        stage_inout: str | None = None,
+        hyperthreads: bool | None = False,
+        dry_run: bool | None = False,
+        wait_on_start: bool | None = False,
+        extra_args: str | None = None,
+        auto_assign_gpus: bool | None = False,
+        match_procs_to_gpus: bool | None = False,
+        env_script: str | None = None,
+        mpi_runner_type: str | dict | None = None,
     ) -> Task:
         """Creates a new task, and either executes or schedules execution.