diff --git a/libensemble/executors/__init__.py b/libensemble/executors/__init__.py index 563fa3352..13c7d851d 100644 --- a/libensemble/executors/__init__.py +++ b/libensemble/executors/__init__.py @@ -1,4 +1,11 @@ from libensemble.executors.executor import Executor from libensemble.executors.mpi_executor import MPIExecutor -__all__ = ["Executor", "MPIExecutor"] +# FluxExecutor is optional - requires flux-core Python bindings +try: + from libensemble.executors.flux_executor import FluxExecutor # noqa: F401 + + __all__ = ["Executor", "MPIExecutor", "FluxExecutor"] +except ImportError: + # flux-core not available - FluxExecutor won't be importable + __all__ = ["Executor", "MPIExecutor"] diff --git a/libensemble/executors/flux_executor.py b/libensemble/executors/flux_executor.py new file mode 100644 index 000000000..8185c9ae8 --- /dev/null +++ b/libensemble/executors/flux_executor.py @@ -0,0 +1,482 @@ +""" +This module provides a native Flux executor using the Flux Python API. + +The FluxExecutor submits jobs directly to Flux using its Python bindings, +rather than wrapping `flux run` as a subprocess. This provides better +integration with Flux's job lifecycle management and is particularly +useful when running inside containers where MPI runners may not be available. + +Usage:: + + from libensemble.executors.flux_executor import FluxExecutor + + exctr = FluxExecutor() + exctr.register_app(full_path="/path/to/my_app.x", app_name="my_app") + + # In your sim function: + task = exctr.submit(app_name="my_app", num_procs=4, num_nodes=1) + task.wait() + +Requirements: + - flux-core Python bindings must be installed + - Must be running inside a Flux instance or provide a Flux URI + +Notes: + Flux handles are not thread-safe. FluxExecutor is best suited for + process-based libEnsemble runs, such as multiprocessing or MPI workers. + The executor reconnects lazily after process serialization so each worker + process uses its own Flux handle. +""" + +import logging +import os +import shlex +import time + +from libensemble.executors.executor import ( + Application, + Executor, + ExecutorException, + Task, + jassert, +) + +logger = logging.getLogger(__name__) + +# Try to import flux - it's optional +try: + import flux + import flux.job + from flux.job import JobspecV1 + + FLUX_AVAILABLE = True +except ImportError: + FLUX_AVAILABLE = False + flux = None + JobspecV1 = None + + +class FluxTask(Task): + """ + Task subclass for Flux jobs using native Flux Python API. + + Overrides poll() and kill() to use Flux job management instead + of subprocess operations. + """ + + def __init__( + self, + app=None, + app_args=None, + workdir=None, + stdout=None, + stderr=None, + workerid=None, + dry_run=False, + ) -> None: + super().__init__(app, app_args, workdir, stdout, stderr, workerid, dry_run) + self.flux_handle = None + self.flux_jobid = None + self.flux_future = None + + def reset(self) -> None: + super().reset() + self.flux_jobid = None + self.flux_future = None + + def _check_poll(self) -> bool: + """Check whether polling this task makes sense.""" + jassert( + self.flux_jobid is not None, + f"task {self.name} has no Flux job ID - check task has been launched", + ) + if self.finished: + logger.debug(f"Polled task {self.name} has already finished. Not re-polling. Status is {self.state}") + return False + return True + + def poll(self) -> None: + """Polls and updates the status attributes of the task using Flux job state.""" + if self.dry_run: + self._set_complete() + return + + if not self._check_poll(): + return + + try: + info = flux.job.get_job(self.flux_handle, self.flux_jobid) + jassert(info is not None, f"Flux job {self.flux_jobid} was not found") + state = str(info.get("state", "UNKNOWN")).upper() + + # Map Flux states to libEnsemble states + # Flux states: DEPEND, PRIORITY, SCHED, RUN, CLEANUP, INACTIVE + if state in ("DEPEND", "PRIORITY", "SCHED"): + self.state = "WAITING" + elif state == "RUN": + self.state = "RUNNING" + self.runtime = self.timer.elapsed + elif state in ("CLEANUP", "INACTIVE"): + # Job has finished - check if successful + self._handle_completion(info) + else: + self.state = "UNKNOWN" + self.runtime = self.timer.elapsed + + except Exception as e: + logger.warning(f"Error polling Flux job {self.flux_jobid}: {e}") + self.state = "UNKNOWN" + self.runtime = self.timer.elapsed + + def _handle_completion(self, info: dict) -> None: + """Handle job completion and determine success/failure.""" + self.finished = True + self.calc_task_timing() + + result = str(info.get("result", "")).upper() + self.errcode = info.get("returncode", 1) + self.success = result == "COMPLETED" or self.errcode == 0 + self.state = "FINISHED" if self.success else "FAILED" + if self.success: + self.errcode = 0 + + logger.info(f"Task {self.name} finished with state {self.state} (result={result})") + + def _handle_result(self, info) -> None: + """Handle a terminal Flux JobInfo object returned by flux.job.result().""" + result = str(getattr(info, "result", "")).upper() + returncode = getattr(info, "returncode", 1) + if returncode == "": + returncode = 0 if result == "COMPLETED" else 1 + + self.errcode = returncode + self.finished = True + self.calc_task_timing() + self.success = result == "COMPLETED" or self.errcode == 0 + self.state = "FINISHED" if self.success else "FAILED" + if self.success: + self.errcode = 0 + + logger.info(f"Task {self.name} finished with state {self.state} (result={result})") + + def _set_complete(self) -> None: + """Set task as complete (used for dry_run).""" + self.finished = True + if self.dry_run: + self.success = True + self.state = "FINISHED" + else: + self.calc_task_timing() + self.success = self.errcode == 0 + self.state = "FINISHED" if self.success else "FAILED" + logger.info(f"Task {self.name} finished with errcode {self.errcode} ({self.state})") + + def wait(self, timeout: float | None = None) -> None: + """Waits on completion of the Flux job or raises TimeoutExpired exception.""" + from libensemble.executors.executor import TimeoutExpired + + if self.dry_run: + self._set_complete() + return + + if not self._check_poll(): + return + + if timeout is not None: + start_time = time.time() + while True: + self.poll() + if self.finished: + return + + if time.time() - start_time >= timeout: + raise TimeoutExpired(self.name, timeout) + + time.sleep(0.1) + + try: + info = flux.job.result(self.flux_handle, self.flux_jobid) + self._handle_result(info) + except Exception as e: + logger.warning(f"Error waiting for Flux job {self.flux_jobid}: {e}") + self.state = "UNKNOWN" + self.runtime = self.timer.elapsed + raise + + def kill(self, wait_time: int | None = 60) -> None: + """Kills/cancels the Flux job. + + Parameters + ---------- + wait_time: int, Optional + Time in seconds to wait for cancellation. + Note: Flux handles job cancellation internally. + """ + self.poll() + if self.dry_run: + return + + if self.finished: + logger.warning(f"Trying to kill task that is no longer running. Task {self.name}: Status is {self.state}") + return + + logger.info(f"Canceling Flux job {self.flux_jobid} for task {self.name}") + + try: + flux.job.cancel(self.flux_handle, self.flux_jobid) + except Exception as e: + logger.warning(f"Error canceling Flux job {self.flux_jobid}: {e}") + return + + if wait_time: + deadline = time.time() + min(wait_time, 5) # Don't wait too long + while time.time() < deadline: + self.poll() + if self.finished: + break + time.sleep(0.1) + + self.state = "USER_KILLED" + self.finished = True + self.success = False + self.calc_task_timing() + + +class FluxExecutor(Executor): + """ + Native Flux executor using the Flux Python API. + + This executor submits jobs directly to Flux rather than wrapping + `flux run` as a subprocess. It provides better integration with + Flux's job lifecycle and is suitable for container-based workflows. + + Parameters + ---------- + uri: str, Optional + Flux instance URI. If omitted, ``flux.Flux()`` connects to the nearest + enclosing Flux instance discovered by Flux. + + Raises + ------ + ExecutorException + If flux Python bindings are not available or connecting to Flux fails. + + Example + ------- + :: + + from libensemble.executors.flux_executor import FluxExecutor + + exctr = FluxExecutor() + exctr.register_app(full_path="/path/to/sim.x", app_name="sim") + + # In sim function: + task = exctr.submit(app_name="sim", num_procs=4) + task.wait() + """ + + def __init__(self, uri: str | None = None) -> None: + """Instantiate a new FluxExecutor instance.""" + if not FLUX_AVAILABLE: + raise ExecutorException( + "Flux Python bindings not available. " + "Install flux-core or use MPIExecutor with mpi_runner='flux' instead." + ) + + super().__init__() + + self.resources = None + self.platform_info: dict = {} + self.uri = uri + self.flux_handle = None + + def __getstate__(self): + """Avoid sharing non-thread-safe Flux handles across worker processes.""" + state = self.__dict__.copy() + state["flux_handle"] = None + return state + + def _get_flux_handle(self): + """Return a Flux handle for this process, opening it lazily if needed.""" + if self.flux_handle is None: + try: + self.flux_handle = flux.Flux(self.uri) if self.uri is not None else flux.Flux() + except Exception as e: + raise ExecutorException(f"Failed to connect to Flux instance: {e}") + return self.flux_handle + + def submit( + self, + 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, + app_args: str | None = None, + stdout: str | None = None, + stderr: str | None = None, + dry_run: bool = False, + wait_on_start: bool | int = False, + extra_args: str | None = None, + ) -> FluxTask: + """Submit a job to Flux. + + Returns :class:`FluxTask` object. + + Parameters + ---------- + calc_type: str, Optional + The calculation type: 'sim' or 'gen' + + app_name: str, Optional + The application name. + + num_procs: int, Optional + The total number of processes (MPI ranks) + + num_nodes: int, Optional + The number of nodes + + procs_per_node: int, Optional + The processes per node + + num_gpus: int, Optional + The total number of GPUs + + app_args: str, Optional + Application arguments + + stdout: str, Optional + Standard output filename + + stderr: str, Optional + Standard error filename + + dry_run: bool, Optional + If True, don't actually submit the job + + wait_on_start: bool or int, Optional + Whether to wait for job to start running. If an integer N is supplied, + wait at most N seconds. + + extra_args: str, Optional + Additional arguments (currently not used for native Flux) + + Returns + ------- + task: FluxTask + The submitted task object + """ + app: Application | None = None + if app_name is not None: + app = self.get_app(app_name) + elif calc_type is not None: + app = self.default_app(calc_type) + else: + raise ExecutorException("Either app_name or calc_type must be set") + + assert app is not None + + default_workdir = os.getcwd() + task = FluxTask(app, app_args, default_workdir, stdout, stderr, self.workerID, dry_run) + if not dry_run: + task.flux_handle = self._get_flux_handle() + + if not dry_run: + self._check_app_exists(task.app) + + if extra_args: + raise ExecutorException("extra_args is not supported by FluxExecutor") + + num_procs = num_procs or 1 + if num_nodes is None: + if procs_per_node is not None: + if num_procs % procs_per_node != 0: + raise ExecutorException("num_procs must be divisible by procs_per_node for FluxExecutor") + num_nodes = num_procs // procs_per_node + elif procs_per_node is not None and num_procs != num_nodes * procs_per_node: + raise ExecutorException("num_procs must equal num_nodes * procs_per_node for FluxExecutor") + + command = shlex.split(task.app.app_cmd) + if task.app_args: + command.extend(shlex.split(task.app_args)) + + command = self._set_sim_dir_env(task, command) + task.runline = " ".join(command) + + if dry_run: + logger.info(f"Test (No submit) Command: {task.runline}") + logger.info(f" num_procs={num_procs}, num_nodes={num_nodes}, procs_per_node={procs_per_node}") + task._set_complete() + else: + # Create Flux jobspec + try: + gpus_per_task = None + if num_gpus is not None: + if num_gpus < 0: + raise ExecutorException("num_gpus must be non-negative") + if num_gpus and num_gpus % num_procs != 0: + raise ExecutorException("num_gpus must be divisible by num_procs for FluxExecutor") + gpus_per_task = num_gpus // num_procs if num_gpus else 0 + + environment = dict(os.environ) + environment.update(task.env) + + jobspec = JobspecV1.from_command( + command, + num_tasks=num_procs, + num_nodes=num_nodes, + cores_per_task=1, + gpus_per_task=gpus_per_task, + cwd=task.workdir, + environment=environment, + output=os.path.join(task.workdir, task.stdout), + error=os.path.join(task.workdir, task.stderr), + ) + if gpus_per_task: + jobspec.setattr_shell_option("gpu-affinity", "per-task") + + logger.info(f"Submitting Flux job for task {task.name}: {task.runline}") + task.flux_future = flux.job.submit_async(task.flux_handle, jobspec) + task.flux_jobid = task.flux_future.get_id() if task.flux_future else None + logger.info(f"Task {task.name} submitted with Flux job ID {task.flux_jobid}") + + task.timer.start() + task.submit_time = task.timer.tstart + + if wait_on_start: + timeout = ( + wait_on_start + if isinstance(wait_on_start, int) and not isinstance(wait_on_start, bool) + else 60.0 + ) + self._wait_on_start(task, timeout) + + except Exception as e: + logger.error(f"Failed to submit Flux job: {e}") + task.state = "FAILED_TO_START" + task.finished = True + raise ExecutorException(f"Failed to submit Flux job: {e}") + + self.list_of_tasks.append(task) + return task + + def _wait_on_start(self, task: FluxTask, timeout: float = 60.0) -> None: + """Wait for a task to start running.""" + start = time.time() + task.timer.start() + task.submit_time = task.timer.tstart + + while task.state in ("CREATED", "WAITING"): + time.sleep(0.1) + task.poll() + if time.time() - start > timeout: + logger.warning(f"Timeout waiting for task {task.name} to start") + break + + if not task.finished: + task.timer.start() + task.submit_time = task.timer.tstart + + logger.debug(f"Task {task.name} polled as {task.state} after {time.time() - start:.2f} seconds") diff --git a/libensemble/tests/unit_tests/test_ensemble.py b/libensemble/tests/unit_tests/test_ensemble.py index 5e5e9314f..7d3ea2d30 100644 --- a/libensemble/tests/unit_tests/test_ensemble.py +++ b/libensemble/tests/unit_tests/test_ensemble.py @@ -284,7 +284,7 @@ def test_gen_specs_vocs_populates_user_bounds(): vocs = VOCS( variables={"x0": [-3, 3], "x1": [-2, 2], "x2": [-1, 1], "x3": [-1, 1]}, - objectives={"f": "EXPLORE"}, + objectives={"f": "EXPLORE"} ) gs = GenSpecs(vocs=vocs) assert "lb" in gs.user, "lb should be populated in user from VOCS" diff --git a/libensemble/tests/unit_tests/test_flux.py b/libensemble/tests/unit_tests/test_flux.py index b5f7c0e01..9d624c5b3 100644 --- a/libensemble/tests/unit_tests/test_flux.py +++ b/libensemble/tests/unit_tests/test_flux.py @@ -6,6 +6,7 @@ - Flux nodelist parsing (via slurm-style bracket notation) - Flux MPI variant detection - FluxAllocation platform configuration +- FluxExecutor (when flux bindings available) """ import os @@ -15,6 +16,8 @@ import pytest +from libensemble.executors import flux_executor +from libensemble.executors.executor import TimeoutExpired from libensemble.executors.mpi_runner import FLUX_MPIRunner, MPIRunner from libensemble.resources.env_resources import EnvResources from libensemble.resources.platforms import FluxAllocation, Known_platforms @@ -317,6 +320,687 @@ class MockCls: check_mpi_runner_type(MockCls, "invalid_runner") +# ======================================================================================== +# Tests for FluxExecutor (conditional on flux availability) +# ======================================================================================== + + +def test_flux_executor_import_without_flux(): + """Test FluxExecutor handles missing flux gracefully""" + # This test just verifies the module can be imported + # even when flux is not available + try: + from libensemble.executors import flux_executor + + # FLUX_AVAILABLE should be False if flux not installed + # This is fine - we just want to ensure import doesn't crash + assert hasattr(flux_executor, "FLUX_AVAILABLE") + except ImportError: + pytest.skip("flux_executor module not available") + + +def test_flux_executor_connects_lazily_with_default_or_explicit_uri(): + """FluxExecutor should defer Flux connection and support explicit URIs.""" + try: + from libensemble.executors.flux_executor import FLUX_AVAILABLE, FluxExecutor + + if not FLUX_AVAILABLE: + pytest.skip("Flux Python bindings not available") + + with mock.patch.object(flux_executor.flux, "Flux", return_value="default-handle") as mock_flux: + executor = FluxExecutor() + assert executor.flux_handle is None + assert executor._get_flux_handle() == "default-handle" + mock_flux.assert_called_once_with() + + with mock.patch.object(flux_executor.flux, "Flux", return_value="uri-handle") as mock_flux: + executor = FluxExecutor(uri="local:///tmp/flux-uri") + assert executor._get_flux_handle() == "uri-handle" + mock_flux.assert_called_once_with("local:///tmp/flux-uri") + + except ImportError: + pytest.skip("flux_executor module not available") + + +def test_flux_task_poll_uses_get_job(): + """Test FluxTask polls using Flux's get_job helper""" + if not flux_executor.FLUX_AVAILABLE: + pytest.skip("Flux Python bindings not available") + + task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=False, + ) + task.flux_handle = object() + task.flux_jobid = 123 + task.timer.start() + task.submit_time = task.timer.tstart + + with mock.patch.object(flux_executor.flux.job, "get_job", return_value={"state": "RUN"}) as mock_get_job: + task.poll() + + mock_get_job.assert_called_once_with(task.flux_handle, task.flux_jobid) + assert task.state == "RUNNING" + + +def _make_uninitialized_flux_executor(*, worker_id: int = 7): + executor = object.__new__(flux_executor.FluxExecutor) + executor.flux_handle = object() + executor.platform_info = {} + executor.uri = None + executor.workerID = worker_id + executor.list_of_tasks = [] + executor.apps = {} + executor.default_apps = {"sim": None, "gen": None} + executor.base_dir = os.getcwd() + return executor + + +def test_flux_executor_submit_builds_jobspec_with_environment_and_gpus(): + """Test FluxExecutor submit passes environment and GPU resources via jobspec""" + + executor = _make_uninitialized_flux_executor(worker_id=7) + + app = SimpleNamespace( + name="sim", full_path="/path/to/sim.x", app_cmd="fluxwrap /path/to/sim.x", precedent="fluxwrap" + ) + executor.get_app = lambda app_name: app + executor.default_app = lambda calc_type: app + executor._check_app_exists = lambda app_obj: None + + old_env = os.environ.get("TEST_FLUX_ENV") + os.environ["TEST_FLUX_ENV"] = "present" + + jobspec = SimpleNamespace(stdout=None, stderr=None) + submit_calls = [] + submit_future = SimpleNamespace(get_id=mock.Mock(return_value=42)) + + def fake_from_command(command, **kwargs): + submit_calls.append((command, kwargs)) + jobspec.cwd = kwargs.get("cwd") + jobspec.environment = kwargs.get("environment") + jobspec.setattr_shell_option = mock.Mock() + return jobspec + + try: + with ( + mock.patch.object(flux_executor, "JobspecV1", create=True, autospec=False) as _patched_jobspecV1, + mock.patch.object(flux_executor, "flux", create=True) as _patched_flux, # noqa: F841 + mock.patch.object(flux_executor.flux, "job", create=True) as _patched_job, # noqa: F841 + mock.patch.object(flux_executor.flux.job, "submit_async", return_value=submit_future) as mock_submit_async, + ): + _patched_jobspecV1.from_command.side_effect = fake_from_command + task = executor.submit(app_name="sim", num_procs=4, num_nodes=2, num_gpus=4, app_args="--flag value") + finally: + if old_env is None: + del os.environ["TEST_FLUX_ENV"] + else: + os.environ["TEST_FLUX_ENV"] = old_env + + command, kwargs = submit_calls[0] + assert command[:2] == ["fluxwrap", "/path/to/sim.x"] + assert command[-2:] == ["--flag", "value"] + assert kwargs["num_tasks"] == 4 + assert kwargs["num_nodes"] == 2 + assert kwargs["gpus_per_task"] == 1 + assert kwargs["environment"]["TEST_FLUX_ENV"] == "present" + assert kwargs["environment"]["LIBENSEMBLE_SIM_DIR"] == "." + assert kwargs["output"] == os.path.join(task.workdir, task.stdout) + assert kwargs["error"] == os.path.join(task.workdir, task.stderr) + mock_submit_async.assert_called_once_with(executor.flux_handle, jobspec) + submit_future.get_id.assert_called_once_with() + jobspec.setattr_shell_option.assert_called_once_with("gpu-affinity", "per-task") + assert task.flux_jobid == 42 + assert task.flux_future is submit_future + + +def test_flux_executor_submit_dry_run_marks_task_complete_and_does_not_submit_job(): + """Dry-run should not call JobspecV1.from_command or flux.job.submit_async.""" + executor = _make_uninitialized_flux_executor(worker_id=7) + + app = SimpleNamespace( + name="sim", + full_path="/path/to/sim.x", + app_cmd="fluxwrap /path/to/sim.x", + precedent="fluxwrap", + ) + + executor.get_app = lambda app_name: app + executor.default_app = lambda calc_type: app + executor._check_app_exists = mock.Mock() + + mock_from_command = mock.Mock() + + with ( + mock.patch.object(flux_executor, "JobspecV1", create=True, autospec=False) as _patched_jobspecV1, + mock.patch.object(flux_executor, "flux", create=True) as _patched_flux, # noqa: F841 + mock.patch.object(flux_executor.flux, "job", create=True) as _patched_job, # noqa: F841 + mock.patch.object(flux_executor.flux.job, "submit_async") as mock_submit, + ): + _patched_jobspecV1.from_command = mock_from_command + task = executor.submit( + app_name="sim", + num_procs=2, + num_nodes=1, + num_gpus=0, + app_args="--flag", + stdout="out.txt", + stderr="err.txt", + dry_run=True, + ) + + mock_from_command.assert_not_called() + mock_submit.assert_not_called() + executor._check_app_exists.assert_not_called() + assert task.finished is True + assert task.state == "FINISHED" + assert task.success is True + assert task.runline is not None + assert len(executor.list_of_tasks) == 1 + + +def test_flux_executor_submit_wait_on_start_invokes_waiter(): + """wait_on_start should call FluxExecutor._wait_on_start when not in dry_run.""" + executor = _make_uninitialized_flux_executor(worker_id=7) + + app = SimpleNamespace( + name="sim", + full_path="/path/to/sim.x", + app_cmd="fluxwrap /path/to/sim.x", + precedent="fluxwrap", + ) + + executor.get_app = lambda app_name: app + executor.default_app = lambda calc_type: app + executor._check_app_exists = lambda app_obj: None + + jobspec = SimpleNamespace(stdout=None, stderr=None) + submit_future = SimpleNamespace(get_id=mock.Mock(return_value=123)) + + with ( + mock.patch.object(flux_executor, "JobspecV1", create=True, autospec=False) as _patched_jobspecV1, + mock.patch.object(flux_executor, "flux", create=True) as _patched_flux, # noqa: F841 + mock.patch.object(flux_executor.flux, "job", create=True) as _patched_job, # noqa: F841 + mock.patch.object(flux_executor.flux.job, "submit_async", return_value=submit_future), + mock.patch.object(executor, "_wait_on_start") as mock_wait_on_start, + ): + _patched_jobspecV1.from_command.return_value = jobspec + task = executor.submit( + app_name="sim", + num_procs=2, + num_nodes=1, + num_gpus=None, + wait_on_start=7, + ) + + mock_wait_on_start.assert_called_once_with(task, 7) + assert task.flux_jobid == 123 + + +def test_flux_executor_submit_validation_errors(): + executor = _make_uninitialized_flux_executor(worker_id=7) + + app = SimpleNamespace( + name="sim", + full_path="/path/to/sim.x", + app_cmd="fluxwrap /path/to/sim.x", + precedent="fluxwrap", + ) + + executor.get_app = lambda app_name: app + executor.default_app = lambda calc_type: app + executor._check_app_exists = lambda app_obj: None + + # Missing app_name + calc_type + with pytest.raises(Exception): + executor.submit() + + # extra_args unsupported + with pytest.raises(Exception, match="extra_args"): + executor.submit(app_name="sim", num_procs=1, num_nodes=1, extra_args="--x") + + # num_gpus negative + with pytest.raises(Exception, match="num_gpus must be non-negative"): + executor.submit(app_name="sim", num_procs=2, num_nodes=1, num_gpus=-1) + + # num_gpus not divisible by num_procs + with pytest.raises(Exception, match="num_gpus must be divisible by num_procs"): + executor.submit(app_name="sim", num_procs=3, num_nodes=1, num_gpus=2) + + # procs_per_node divides num_procs -> num_nodes inferred + with ( + mock.patch.object(flux_executor, "JobspecV1", create=True, autospec=False) as _patched_jobspecV1, + mock.patch.object(flux_executor, "flux", create=True) as _patched_flux, # noqa: F841 + mock.patch.object(flux_executor.flux, "job", create=True) as _patched_job, # noqa: F841 + mock.patch.object( + flux_executor.flux.job, "submit_async", return_value=SimpleNamespace(get_id=mock.Mock(return_value=1)) + ), + ): + _patched_jobspecV1.from_command.return_value = SimpleNamespace(stdout=None, stderr=None) + executor.submit(app_name="sim", num_procs=4, procs_per_node=2) + + # num_procs must be divisible by procs_per_node + with pytest.raises(Exception, match="divisible by procs_per_node"): + executor.submit(app_name="sim", num_procs=3, procs_per_node=2) + + # num_procs must equal num_nodes * procs_per_node + with pytest.raises(Exception, match=r"num_procs must equal num_nodes \* procs_per_node"): + executor.submit(app_name="sim", num_procs=5, num_nodes=2, procs_per_node=3) + + +def test_flux_executor_submit_error_from_flux_job_submit_sets_failed_to_start(): + executor = _make_uninitialized_flux_executor(worker_id=7) + + app = SimpleNamespace( + name="sim", + full_path="/path/to/sim.x", + app_cmd="fluxwrap /path/to/sim.x", + precedent="fluxwrap", + ) + + executor.get_app = lambda app_name: app + executor.default_app = lambda calc_type: app + executor._check_app_exists = lambda app_obj: None + + jobspec = SimpleNamespace(stdout=None, stderr=None) + + with ( + mock.patch.object(flux_executor, "JobspecV1", create=True, autospec=False) as _patched_jobspecV1, + mock.patch.object(flux_executor, "flux", create=True) as _patched_flux, # noqa: F841 + mock.patch.object(flux_executor.flux, "job", create=True) as _patched_job, # noqa: F841 + mock.patch.object(flux_executor.flux.job, "submit_async", side_effect=RuntimeError("boom")), + ): + _patched_jobspecV1.from_command.return_value = jobspec + with pytest.raises(Exception, match="Failed to submit Flux job"): + executor.submit(app_name="sim", num_procs=2, num_nodes=1) + + # Submit failed; the task is created, but it may or may not be appended + # depending on where the exception is raised. Ensure the exception type is correct. + assert executor.list_of_tasks == [] + + +def test_flux_executor_submit_stdout_stderr_are_placed_under_workdir(): + executor = _make_uninitialized_flux_executor(worker_id=7) + + app = SimpleNamespace( + name="sim", + full_path="/path/to/sim.x", + app_cmd="fluxwrap /path/to/sim.x", + precedent="fluxwrap", + ) + + executor.get_app = lambda app_name: app + executor.default_app = lambda calc_type: app + executor._check_app_exists = lambda app_obj: None + + submit_calls = [] + + def fake_from_command(command, **kwargs): + submit_calls.append(kwargs) + return SimpleNamespace(stdout=None, stderr=None) + + with ( + mock.patch.object(flux_executor, "JobspecV1", create=True, autospec=False) as _patched_jobspecV1, + mock.patch.object(flux_executor, "flux", create=True) as _patched_flux, # noqa: F841 + mock.patch.object(flux_executor.flux, "job", create=True) as _patched_job, # noqa: F841 + mock.patch.object( + flux_executor.flux.job, "submit_async", return_value=SimpleNamespace(get_id=mock.Mock(return_value=1)) + ), + ): + _patched_jobspecV1.from_command.side_effect = fake_from_command + task = executor.submit( + app_name="sim", + num_procs=2, + num_nodes=1, + stdout="my_stdout.txt", + stderr="my_stderr.txt", + ) + + assert task.flux_jobid == 1 + assert submit_calls[0]["output"].endswith(os.path.join(task.workdir, "my_stdout.txt")) + assert submit_calls[0]["error"].endswith(os.path.join(task.workdir, "my_stderr.txt")) + + +def test_flux_executor_submit_sets_gpu_affinity_only_when_num_gpus_nonzero(): + executor = _make_uninitialized_flux_executor(worker_id=7) + + app = SimpleNamespace( + name="sim", + full_path="/path/to/sim.x", + app_cmd="fluxwrap /path/to/sim.x", + precedent="fluxwrap", + ) + + executor.get_app = lambda app_name: app + executor.default_app = lambda calc_type: app + executor._check_app_exists = lambda app_obj: None + + jobspec = SimpleNamespace(stdout=None, stderr=None) + jobspec.setattr_shell_option = mock.Mock() + + with ( + mock.patch.object(flux_executor, "JobspecV1", create=True, autospec=False) as _patched_jobspecV1, + mock.patch.object(flux_executor, "flux", create=True) as _patched_flux, # noqa: F841 + mock.patch.object(flux_executor.flux, "job", create=True) as _patched_job, # noqa: F841 + mock.patch.object( + flux_executor.flux.job, "submit_async", return_value=SimpleNamespace(get_id=mock.Mock(return_value=1)) + ), + ): + _patched_jobspecV1.from_command.return_value = jobspec + executor.submit(app_name="sim", num_procs=4, num_nodes=1, num_gpus=0) + + jobspec.setattr_shell_option.assert_not_called() + + +def test_flux_executor_getstate_drops_flux_handle(): + """FluxExecutor should not serialize an open Flux handle into worker processes.""" + with mock.patch.object(flux_executor, "FLUX_AVAILABLE", True): + executor = flux_executor.FluxExecutor(uri="local:///tmp/flux-test") + + executor.flux_handle = "flux-handle" + state = executor.__getstate__() + assert state["flux_handle"] is None + assert executor.flux_handle == "flux-handle" + + +def test_flux_executor_wait_on_start_polls_until_running(): + """Test FluxExecutor waits for a FluxTask to leave the startup states.""" + executor = object.__new__(flux_executor.FluxExecutor) + task = SimpleNamespace( + name="flux-task", + state="CREATED", + finished=False, + timer=SimpleNamespace(tstart=None, start=mock.Mock(side_effect=lambda: setattr(task.timer, "tstart", 1.23))), + submit_time=None, + ) + + def poll_side_effect(): + task.state = "RUNNING" + + task.poll = mock.Mock(side_effect=poll_side_effect) + + with mock.patch.object(flux_executor.time, "sleep"): + executor._wait_on_start(task, timeout=0.5) + + assert task.poll.call_count == 1 + assert task.state == "RUNNING" + assert task.timer.start.call_count == 2 + assert task.submit_time == 1.23 + + +def test_flux_task_poll_maps_completion_waiting_and_unknown_states(): + """Test FluxTask poll maps Flux job states to libEnsemble states.""" + task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=False, + ) + task.flux_handle = object() + task.flux_jobid = 123 + task.timer.start() + task.submit_time = task.timer.tstart + fake_flux = SimpleNamespace(job=SimpleNamespace(get_job=mock.Mock())) + + with mock.patch.object(flux_executor, "flux", fake_flux): + fake_flux.job.get_job.return_value = {"state": "SCHED"} + task.poll() + assert task.state == "WAITING" + assert not task.finished + + with mock.patch.object(task, "_handle_completion") as mock_handle_completion: + fake_flux.job.get_job.return_value = {"state": "INACTIVE"} + task.poll() + mock_handle_completion.assert_called_once_with({"state": "INACTIVE"}) + + fake_flux.job.get_job.return_value = {"state": "MYSTERY"} + task.finished = False + task.poll() + + assert task.state == "UNKNOWN" + + +def test_flux_task_handle_completion_success_and_failure(): + """Test FluxTask completion handling sets success, state, and errcode.""" + success_task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=False, + ) + success_task.timer.start() + success_task.submit_time = success_task.timer.tstart + success_task._handle_completion({"state": "INACTIVE", "result": "COMPLETED", "returncode": 0}) + assert success_task.finished is True + assert success_task.success is True + assert success_task.state == "FINISHED" + assert success_task.errcode == 0 + + failed_task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=False, + ) + failed_task.timer.start() + failed_task.submit_time = failed_task.timer.tstart + failed_task._handle_completion({"state": "INACTIVE", "result": "FAILED", "returncode": 7}) + assert failed_task.finished is True + assert failed_task.success is False + assert failed_task.state == "FAILED" + assert failed_task.errcode == 7 + + +def test_flux_task_set_complete_handles_dry_run_and_return_codes(): + """Test FluxTask _set_complete for dry-run and non-dry-run tasks.""" + dry_run_task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=True, + ) + dry_run_task._set_complete() + assert dry_run_task.finished is True + assert dry_run_task.success is True + assert dry_run_task.state == "FINISHED" + + finished_task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=False, + ) + finished_task.errcode = 3 + finished_task.timer.start() + finished_task.submit_time = finished_task.timer.tstart + finished_task._set_complete() + assert finished_task.finished is True + assert finished_task.success is False + assert finished_task.state == "FAILED" + + # cover waiting on a task that completes before timeout + task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=False, + ) + task._set_complete() + task.flux_jobid = 123 + task.wait(timeout=10) + task.kill() + + +def test_flux_task_dry_run_exception_and_kill(): + """Test FluxTask dry run exception attributes.""" + task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=True, + ) + task.wait() + assert task.finished is True + assert task.success is True + assert task.state == "FINISHED" + task.kill() + task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=True, + ) + task.poll() + assert task.finished is True + assert task.success is True + assert task.state == "FINISHED" + + +def test_flux_task_wait_completes_and_times_out(): + """Test FluxTask wait completes after polling and raises on timeout.""" + task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=False, + ) + task.flux_handle = object() + task.flux_jobid = 123 + + def complete_on_second_poll(): + complete_on_second_poll.calls += 1 + if complete_on_second_poll.calls == 1: + task.state = "RUNNING" + else: + task.finished = True + task.state = "FINISHED" + + complete_on_second_poll.calls = 0 + task.poll = mock.Mock(side_effect=complete_on_second_poll) + + with mock.patch.object(flux_executor.time, "sleep"): + task.wait(timeout=1.0) + + assert task.finished is True + assert task.state == "FINISHED" + assert task.poll.call_count == 2 + + result_task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=False, + ) + result_task.flux_handle = object() + result_task.flux_jobid = 321 + result_task.timer.start() + result_task.submit_time = result_task.timer.tstart + result_info = SimpleNamespace(result="COMPLETED", returncode=0) + fake_flux = SimpleNamespace(job=SimpleNamespace(result=mock.Mock(return_value=result_info))) + + with mock.patch.object(flux_executor, "flux", fake_flux): + result_task.wait() + + fake_flux.job.result.assert_called_once_with(result_task.flux_handle, result_task.flux_jobid) + assert result_task.finished is True + assert result_task.state == "FINISHED" + + timeout_task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=False, + ) + timeout_task.flux_handle = object() + timeout_task.flux_jobid = 456 + timeout_task.poll = mock.Mock(side_effect=lambda: setattr(timeout_task, "state", "RUNNING")) + + with ( + mock.patch.object(flux_executor.time, "sleep"), + mock.patch.object(flux_executor.time, "time", side_effect=[0.0, 0.2]), + ): + with pytest.raises(TimeoutExpired): + timeout_task.wait(timeout=0.1) + + +def test_flux_task_kill_cancels_and_marks_user_killed(): + """Test FluxTask kill cancels the job and marks the task as user-killed.""" + task = flux_executor.FluxTask( + app=SimpleNamespace(name="app"), + app_args=None, + workdir=os.getcwd(), + stdout="out.txt", + stderr="err.txt", + workerid=1, + dry_run=False, + ) + task.flux_handle = object() + task.flux_jobid = 789 + task.timer.start() + task.submit_time = task.timer.tstart + + def poll_side_effect(): + if poll_side_effect.calls == 0: + task.state = "RUNNING" + else: + task.finished = True + task.state = "FAILED" + poll_side_effect.calls += 1 + + poll_side_effect.calls = 0 + task.poll = mock.Mock(side_effect=poll_side_effect) + fake_flux = SimpleNamespace(job=SimpleNamespace(cancel=mock.Mock())) + + with ( + mock.patch.object(flux_executor, "flux", fake_flux), + mock.patch.object(flux_executor.time, "sleep"), + mock.patch.object(flux_executor.time, "time", side_effect=[0.0, 0.0, 0.2]), + ): + task.kill(wait_time=1) + + fake_flux.job.cancel.assert_called_once_with(task.flux_handle, task.flux_jobid) + assert task.state == "USER_KILLED" + assert task.finished is True + + # ======================================================================================== # Test runner standalone execution # ======================================================================================== @@ -342,6 +1026,7 @@ class MockCls: # Validator tests test_validator_accepts_flux() test_validator_accepts_all_runners() + test_validator_rejects_invalid() # Platform tests test_flux_allocation_platform() @@ -350,4 +1035,14 @@ class MockCls: # EnvResources tests test_env_resources_flux_env_variable() + # Flux Executor tests + test_flux_executor_getstate_drops_flux_handle() + test_flux_executor_wait_on_start_polls_until_running() + test_flux_task_poll_maps_completion_waiting_and_unknown_states() + test_flux_task_handle_completion_success_and_failure() + test_flux_task_set_complete_handles_dry_run_and_return_codes() + test_flux_task_dry_run_exception_and_kill() + test_flux_task_wait_completes_and_times_out() + test_flux_task_kill_cancels_and_marks_user_killed() + print("All standalone tests passed!") diff --git a/pixi.lock b/pixi.lock index 4ffdd63bc..71990c9d0 100644 --- a/pixi.lock +++ b/pixi.lock @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:17f90f7745889e8d6d423a8aa93c74ef9787036f95086ce5ec38b983de3d36f4 -size 1245374 +oid sha256:94625240a210e869d7239f4709c9db1312145c6d53e876b0a25d70e7e8d87ead +size 1248106 diff --git a/pyproject.toml b/pyproject.toml index 693cd130e..041c10fa6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -173,7 +173,17 @@ globus-compute-sdk = ">=4.10.2,<5" [tool.pixi.feature.py312e.dependencies] globus-compute-sdk = ">=4.10.2,<5" -[tool.pixi.feature.py313e] +[tool.pixi.feature.py312e.target.linux-64.dependencies] +flux-core = ">=0.81.0,<0.82" + +[tool.pixi.feature.py312e.target.linux-64.pypi-dependencies] +flux-python = ">=0.81.0, <0.82" + +[tool.pixi.feature.py313e.target.linux-64.dependencies] +flux-core = ">=0.81.0,<0.82" + +[tool.pixi.feature.py313e.target.linux-64.pypi-dependencies] +flux-python = ">=0.81.0, <0.82" [tool.pixi.feature.py314e]