parsl 2024.2.12__py3-none-any.whl → 2024.2.26__py3-none-any.whl

parsl/tests/test_htex/test_htex.py

@@ -1,18 +1,40 @@
 import pathlib
+from unittest import mock
 
 import pytest
 
 from parsl import curvezmq
 from parsl import HighThroughputExecutor
+from parsl.multiprocessing import ForkProcess
+
+_MOCK_BASE = "parsl.executors.high_throughput.executor"
+
+
+@pytest.fixture
+def encrypted(request: pytest.FixtureRequest):
+    if hasattr(request, "param"):
+        return request.param
+    return True
+
+
+@pytest.fixture
+def htex(encrypted: bool):
+    htex = HighThroughputExecutor(encrypted=encrypted)
+
+    yield htex
+
+    htex.shutdown()
 
 
 @pytest.mark.local
-@pytest.mark.parametrize("encrypted", (True, False))
+@pytest.mark.parametrize("encrypted", (True, False), indirect=True)
 @pytest.mark.parametrize("cert_dir_provided", (True, False))
 def test_htex_start_encrypted(
-    encrypted: bool, cert_dir_provided: bool, tmpd_cwd: pathlib.Path
+    encrypted: bool,
+    cert_dir_provided: bool,
+    htex: HighThroughputExecutor,
+    tmpd_cwd: pathlib.Path,
 ):
-    htex = HighThroughputExecutor(encrypted=encrypted)
     htex.run_dir = str(tmpd_cwd)
     if cert_dir_provided:
         provided_base_dir = tmpd_cwd / "provided"
@@ -44,3 +66,44 @@ def test_htex_start_encrypted(
         assert htex.outgoing_q.zmq_context.cert_dir is None
         assert htex.incoming_q.zmq_context.cert_dir is None
         assert htex.command_client.zmq_context.cert_dir is None
+
+
+@pytest.mark.local
+@pytest.mark.parametrize("started", (True, False))
+@pytest.mark.parametrize("timeout_expires", (True, False))
+@mock.patch(f"{_MOCK_BASE}.logger")
+def test_htex_shutdown(
+    mock_logger: mock.MagicMock,
+    started: bool,
+    timeout_expires: bool,
+    htex: HighThroughputExecutor,
+):
+    mock_ix_proc = mock.Mock(spec=ForkProcess)
+
+    if started:
+        htex.interchange_proc = mock_ix_proc
+        mock_ix_proc.is_alive.return_value = True
+
+    if not timeout_expires:
+        # Simulate termination of the Interchange process
+        def kill_interchange(*args, **kwargs):
+            mock_ix_proc.is_alive.return_value = False
+
+        mock_ix_proc.terminate.side_effect = kill_interchange
+
+    htex.shutdown()
+
+    mock_logs = mock_logger.info.call_args_list
+    if started:
+        assert mock_ix_proc.terminate.called
+        assert mock_ix_proc.join.called
+        assert {"timeout": 10} == mock_ix_proc.join.call_args[1]
+        if timeout_expires:
+            assert "Unable to terminate Interchange" in mock_logs[1][0][0]
+            assert mock_ix_proc.kill.called
+        assert "Attempting" in mock_logs[0][0][0]
+        assert "Finished" in mock_logs[-1][0][0]
+    else:
+        assert not mock_ix_proc.terminate.called
+        assert not mock_ix_proc.join.called
+        assert "has not started" in mock_logs[0][0][0]

parsl/tests/test_monitoring/test_incomplete_futures.py

@@ -0,0 +1,65 @@
+import logging
+import os
+import parsl
+import pytest
+import random
+
+from concurrent.futures import Future
+
+
+@parsl.python_app
+def this_app(inputs=()):
+    return inputs[0]
+
+
+@pytest.mark.local
+def test_future_representation(tmpd_cwd):
+    import sqlalchemy
+    from sqlalchemy import text
+    from parsl.tests.configs.htex_local_alternate import fresh_config
+
+    monitoring_db = str(tmpd_cwd / "monitoring.db")
+    monitoring_url = "sqlite:///" + monitoring_db
+
+    c = fresh_config()
+    c.monitoring.logging_endpoint = monitoring_url
+    c.run_dir = tmpd_cwd
+
+    parsl.load(c)
+
+    # we're going to pass this TOKEN into an app via a pre-requisite Future,
+    # and then expect to see it appear in the monitoring database.
+    TOKEN = random.randint(0, 1000000)
+
+    # make a Future that has no result yet
+    # invoke a task that depends on it
+    # inspect and insert something about the monitoring recorded value of that Future
+    # make the Future complete
+    # inspect and insert something about the monitoring recorded value of that Future
+
+    f1 = Future()
+
+    f2 = this_app(inputs=[f1])
+
+    f1.set_result(TOKEN)
+
+    assert f2.result() == TOKEN
+
+    # this cleanup gives a barrier that allows the monitoring code to store
+    # everything it has in the database - without this, there's a race
+    # condition that the task will not have arrived in the database yet.
+    # A different approach for this test might be to poll the DB for a few
+    # seconds, with the assumption "data will arrive in the DB within
+    # 30 seconds, but probably much sooner".
+    parsl.dfk().cleanup()
+    parsl.clear()
+
+    engine = sqlalchemy.create_engine(monitoring_url)
+    with engine.begin() as connection:
+        result = connection.execute(text("SELECT COUNT(*) FROM task"))
+        (task_count, ) = result.first()
+        assert task_count == 1
+
+        result = connection.execute(text("SELECT task_inputs FROM task"))
+        (task_inputs, ) = result.first()
+        assert task_inputs == "[" + repr(TOKEN) + "]"

parsl/tests/test_mpi_apps/test_bad_mpi_config.py

@@ -0,0 +1,41 @@
+import pytest
+
+from parsl import Config
+from parsl.executors import HighThroughputExecutor
+from parsl.launchers import SrunLauncher, SingleNodeLauncher, SimpleLauncher, AprunLauncher
+from parsl.providers import SlurmProvider
+
+
+@pytest.mark.local
+def test_bad_launcher_with_mpi_mode():
+    """AssertionError if a launcher other than SingleNodeLauncher is supplied"""
+
+    for launcher in [SrunLauncher(), SimpleLauncher(), AprunLauncher()]:
+        with pytest.raises(AssertionError):
+            Config(executors=[
+                HighThroughputExecutor(
+                    enable_mpi_mode=True,
+                    provider=SlurmProvider(launcher=launcher),
+                )
+            ])
+
+
+@pytest.mark.local
+def test_correct_launcher_with_mpi_mode():
+    """Confirm that SingleNodeLauncer works with mpi_mode"""
+
+    config = Config(executors=[
+        HighThroughputExecutor(
+            enable_mpi_mode=True,
+            provider=SlurmProvider(launcher=SingleNodeLauncher()),
+        )
+    ])
+    assert isinstance(config.executors[0].provider.launcher, SingleNodeLauncher)
+
+    config = Config(executors=[
+        HighThroughputExecutor(
+            enable_mpi_mode=True,
+            provider=SlurmProvider(),
+        )
+    ])
+    assert isinstance(config.executors[0].provider.launcher, SingleNodeLauncher)

parsl/tests/test_mpi_apps/test_mpi_mode_disabled.py

@@ -0,0 +1,51 @@
+import logging
+from typing import Dict
+import pytest
+import parsl
+from parsl import python_app
+from parsl.tests.configs.htex_local import fresh_config
+
+EXECUTOR_LABEL = "MPI_TEST"
+
+
+def local_setup():
+    config = fresh_config()
+    config.executors[0].label = EXECUTOR_LABEL
+    config.executors[0].max_workers = 1
+    config.executors[0].enable_mpi_mode = False
+    parsl.load(config)
+
+
+def local_teardown():
+    parsl.dfk().cleanup()
+    parsl.clear()
+
+
+@python_app
+def get_env_vars(parsl_resource_specification: Dict = {}) -> Dict:
+    import os
+
+    parsl_vars = {}
+    for key in os.environ:
+        if key.startswith("PARSL_"):
+            parsl_vars[key] = os.environ[key]
+    return parsl_vars
+
+
+@pytest.mark.local
+def test_only_resource_specs_set():
+    """Confirm that resource_spec env vars are set while launch prefixes are not
+    when enable_mpi_mode = False"""
+    resource_spec = {
+        "num_nodes": 4,
+        "ranks_per_node": 2,
+    }
+
+    future = get_env_vars(parsl_resource_specification=resource_spec)
+
+    result = future.result()
+    assert isinstance(result, Dict)
+    assert "PARSL_DEFAULT_PREFIX" not in result
+    assert "PARSL_SRUN_PREFIX" not in result
+    assert result["PARSL_NUM_NODES"] == str(resource_spec["num_nodes"])
+    assert result["PARSL_RANKS_PER_NODE"] == str(resource_spec["ranks_per_node"])

parsl/tests/test_mpi_apps/test_mpi_mode_enabled.py

@@ -0,0 +1,171 @@
+import logging
+import random
+from typing import Dict
+import pytest
+import parsl
+from parsl import python_app, bash_app
+from parsl.tests.configs.htex_local import fresh_config
+
+import os
+
+EXECUTOR_LABEL = "MPI_TEST"
+
+
+def local_setup():
+    config = fresh_config()
+    config.executors[0].label = EXECUTOR_LABEL
+    config.executors[0].max_workers = 2
+    config.executors[0].enable_mpi_mode = True
+    config.executors[0].mpi_launcher = "mpiexec"
+
+    cwd = os.path.abspath(os.path.dirname(__file__))
+    pbs_nodefile = os.path.join(cwd, "mocks", "pbs_nodefile")
+
+    config.executors[0].provider.worker_init = f"export PBS_NODEFILE={pbs_nodefile}"
+
+    parsl.load(config)
+
+
+def local_teardown():
+    parsl.dfk().cleanup()
+    parsl.clear()
+
+
+@python_app
+def get_env_vars(parsl_resource_specification: Dict = {}) -> Dict:
+    import os
+
+    parsl_vars = {}
+    for key in os.environ:
+        if key.startswith("PARSL_"):
+            parsl_vars[key] = os.environ[key]
+    return parsl_vars
+
+
+@pytest.mark.local
+def test_only_resource_specs_set():
+    """Confirm that resource_spec env vars are set while launch prefixes are not
+    when enable_mpi_mode = False"""
+    resource_spec = {
+        "num_nodes": 2,
+        "ranks_per_node": 2,
+    }
+
+    future = get_env_vars(parsl_resource_specification=resource_spec)
+
+    result = future.result()
+    assert isinstance(result, Dict)
+    logging.warning(f"Got table: {result}")
+    assert "PARSL_MPI_PREFIX" in result
+    assert "PARSL_MPIEXEC_PREFIX" in result
+    assert result["PARSL_MPI_PREFIX"] == result["PARSL_MPIEXEC_PREFIX"]
+    assert result["PARSL_NUM_NODES"] == str(resource_spec["num_nodes"])
+    assert result["PARSL_RANKS_PER_NODE"] == str(resource_spec["ranks_per_node"])
+    assert result["PARSL_NUM_RANKS"] == str(resource_spec["ranks_per_node"] * resource_spec["num_nodes"])
+
+
+@bash_app
+def echo_launch_cmd(
+    parsl_resource_specification: Dict,
+    stdout=parsl.AUTO_LOGNAME,
+    stderr=parsl.AUTO_LOGNAME,
+):
+    return 'echo "$PARSL_MPI_PREFIX hostname"'
+
+
+@pytest.mark.local
+def test_bash_default_prefix_set():
+    """Confirm that resource_spec env vars are set while launch prefixes are not
+    when enable_mpi_mode = False"""
+    resource_spec = {
+        "num_nodes": 2,
+        "ranks_per_node": 2,
+    }
+
+    future = echo_launch_cmd(parsl_resource_specification=resource_spec)
+
+    result = future.result()
+    assert result == 0
+    with open(future.stdout) as f:
+        output = f.readlines()
+        assert output[0].startswith("mpiexec")
+        logging.warning(f"output : {output}")
+
+
+@pytest.mark.local
+def test_bash_multiple_set():
+    """Confirm that multiple apps can run without blocking each other out
+    when enable_mpi_mode = False"""
+    resource_spec = {
+        "num_nodes": 2,
+        "num_ranks": 4,
+    }
+    futures = []
+    for i in range(4):
+        resource_spec["num_nodes"] = i + 1
+        future = echo_launch_cmd(parsl_resource_specification=resource_spec)
+        futures.append(future)
+
+    for future in futures:
+        result = future.result()
+        assert result == 0
+        with open(future.stdout) as f:
+            output = f.readlines()
+            assert output[0].startswith("mpiexec")
+
+
+@bash_app
+def bash_resource_spec(parsl_resource_specification=None, stdout=parsl.AUTO_LOGNAME):
+    total_ranks = (
+        parsl_resource_specification["ranks_per_node"] * parsl_resource_specification["num_nodes"]
+    )
+    return f'echo "{total_ranks}"'
+
+
+@pytest.mark.local
+def test_bash_app_using_resource_spec():
+    resource_spec = {
+        "num_nodes": 2,
+        "ranks_per_node": 2,
+    }
+    future = bash_resource_spec(parsl_resource_specification=resource_spec)
+    assert future.result() == 0
+    with open(future.stdout) as f:
+        output = f.readlines()
+        total_ranks = resource_spec["num_nodes"] * resource_spec["ranks_per_node"]
+        assert int(output[0].strip()) == total_ranks
+
+
+@python_app
+def mock_app(sleep_dur: float = 0.0, parsl_resource_specification: Dict = {}):
+    import os
+    import time
+    time.sleep(sleep_dur)
+
+    total_ranks = int(os.environ["PARSL_NUM_NODES"]) * int(os.environ["PARSL_RANKS_PER_NODE"])
+    nodes = os.environ["PARSL_MPI_NODELIST"].split(',')
+
+    return total_ranks, nodes
+
+
+@pytest.mark.local
+def test_simulated_load(rounds: int = 100):
+
+    node_choices = (1, 2, 4)
+    sleep_choices = (0, 0.01, 0.02, 0.04)
+    ranks_per_node = (4, 8)
+
+    futures = {}
+    for i in range(rounds):
+        resource_spec = {
+            "num_nodes": random.choice(node_choices),
+            "ranks_per_node": random.choice(ranks_per_node),
+        }
+        future = mock_app(sleep_dur=random.choice(sleep_choices),
+                          parsl_resource_specification=resource_spec)
+        futures[future] = resource_spec
+
+    for future in futures:
+        total_ranks, nodes = future.result(timeout=10)
+        assert len(nodes) == futures[future]["num_nodes"]
+        assert total_ranks == futures[future]["num_nodes"] * futures[future]["ranks_per_node"]

parsl/tests/test_mpi_apps/test_mpi_prefix.py

@@ -0,0 +1,71 @@
+import logging
+import pytest
+
+from parsl.executors.high_throughput.mpi_resource_management import Scheduler
+from parsl.executors.high_throughput.mpi_prefix_composer import (
+    compose_srun_launch_cmd,
+    compose_aprun_launch_cmd,
+    compose_mpiexec_launch_cmd,
+    compose_all,
+)
+
+
+resource_spec = {"num_nodes": 2,
+                 "num_ranks": 8,
+                 "ranks_per_node": 4}
+
+
+@pytest.mark.local
+def test_srun_launch_cmd():
+    prefix_name, composed_prefix = compose_srun_launch_cmd(
+        resource_spec=resource_spec, node_hostnames=["node1", "node2"]
+    )
+    assert prefix_name == "PARSL_SRUN_PREFIX"
+    logging.warning(composed_prefix)
+
+    assert "None" not in composed_prefix
+
+
+@pytest.mark.local
+def test_aprun_launch_cmd():
+    prefix_name, composed_prefix = compose_aprun_launch_cmd(
+        resource_spec=resource_spec, node_hostnames=["node1", "node2"]
+    )
+    logging.warning(composed_prefix)
+    assert prefix_name == "PARSL_APRUN_PREFIX"
+    assert "None" not in composed_prefix
+
+
+@pytest.mark.local
+def test_mpiexec_launch_cmd():
+    prefix_name, composed_prefix = compose_mpiexec_launch_cmd(
+        resource_spec=resource_spec, node_hostnames=["node1", "node2"]
+    )
+    logging.warning(composed_prefix)
+    assert prefix_name == "PARSL_MPIEXEC_PREFIX"
+    assert "None" not in composed_prefix
+
+
+@pytest.mark.local
+def test_slurm_launch_cmd():
+    table = compose_all(
+        mpi_launcher="srun",
+        resource_spec=resource_spec,
+        node_hostnames=["NODE001", "NODE002"],
+    )
+
+    assert "PARSL_MPI_PREFIX" in table
+    assert "PARSL_SRUN_PREFIX" in table
+
+
+@pytest.mark.local
+def test_default_launch_cmd():
+    table = compose_all(
+        mpi_launcher="srun",
+        resource_spec=resource_spec,
+        node_hostnames=["NODE001", "NODE002"],
+    )
+
+    assert "PARSL_MPI_PREFIX" in table
+    assert "PARSL_SRUN_PREFIX" in table
+    assert table["PARSL_MPI_PREFIX"] == table["PARSL_SRUN_PREFIX"]

parsl/tests/test_mpi_apps/test_mpi_scheduler.py

@@ -0,0 +1,158 @@
+import logging
+import os
+from unittest import mock
+import pytest
+import pickle
+from parsl.executors.high_throughput.mpi_resource_management import TaskScheduler, MPITaskScheduler
+from parsl.multiprocessing import SpawnContext
+from parsl.serialize import pack_res_spec_apply_message, unpack_res_spec_apply_message
+
+
+@pytest.fixture(autouse=True)
+def set_pbs_nodefile_envvars():
+    cwd = os.path.abspath(os.path.dirname(__file__))
+    pbs_nodefile = os.path.join(cwd, "mocks", "pbs_nodefile.8")
+    with mock.patch.dict(os.environ, {"PBS_NODEFILE": pbs_nodefile}):
+        yield
+
+
+@pytest.mark.local
+def test_NoopScheduler():
+    task_q, result_q = SpawnContext.Queue(), SpawnContext.Queue()
+    scheduler = TaskScheduler(task_q, result_q)
+
+    scheduler.put_task("TaskFoo")
+    assert task_q.get() == "TaskFoo"
+
+    result_q.put("Result1")
+    assert scheduler.get_result(True, 1) == "Result1"
+
+
+@pytest.mark.local
+def test_MPISched_put_task():
+    task_q, result_q = SpawnContext.Queue(), SpawnContext.Queue()
+    scheduler = MPITaskScheduler(task_q, result_q)
+
+    assert scheduler.available_nodes
+    assert len(scheduler.available_nodes) == 8
+    assert scheduler._free_node_counter.value == 8
+
+    mock_task_buffer = pack_res_spec_apply_message("func",
+                                                   "args",
+                                                   "kwargs",
+                                                   resource_specification={"num_nodes": 2,
+                                                                           "ranks_per_node": 2})
+    task_package = {"task_id": 1, "buffer": mock_task_buffer}
+    scheduler.put_task(task_package)
+
+    assert scheduler._free_node_counter.value == 6
+
+
+@pytest.mark.local
+def test_MPISched_get_result():
+    task_q, result_q = SpawnContext.Queue(), SpawnContext.Queue()
+    scheduler = MPITaskScheduler(task_q, result_q)
+
+    assert scheduler.available_nodes
+    assert len(scheduler.available_nodes) == 8
+    assert scheduler._free_node_counter.value == 8
+
+    nodes = [scheduler.nodes_q.get() for _ in range(4)]
+    scheduler._free_node_counter.value = 4
+    scheduler._map_tasks_to_nodes[1] = nodes
+
+    result_package = pickle.dumps({"task_id": 1, "type": "result", "buffer": "Foo"})
+    result_q.put(result_package)
+    result_received = scheduler.get_result(block=True, timeout=1)
+    assert result_received == result_package
+
+    assert scheduler._free_node_counter.value == 8
+
+
+@pytest.mark.local
+def test_MPISched_roundtrip():
+    task_q, result_q = SpawnContext.Queue(), SpawnContext.Queue()
+    scheduler = MPITaskScheduler(task_q, result_q)
+
+    assert scheduler.available_nodes
+    assert len(scheduler.available_nodes) == 8
+
+    for round in range(1, 9):
+        assert scheduler._free_node_counter.value == 8
+
+        mock_task_buffer = pack_res_spec_apply_message("func",
+                                                       "args",
+                                                       "kwargs",
+                                                       resource_specification={"num_nodes": round,
+                                                                               "ranks_per_node": 2})
+        task_package = {"task_id": round, "buffer": mock_task_buffer}
+        scheduler.put_task(task_package)
+
+        assert scheduler._free_node_counter.value == 8 - round
+
+        # Pop in a mock result
+        result_pkl = pickle.dumps({"task_id": round, "type": "result", "buffer": "RESULT BUF"})
+        result_q.put(result_pkl)
+
+        got_result = scheduler.get_result(True, 1)
+        assert got_result == result_pkl
+
+
+@pytest.mark.local
+def test_MPISched_contention():
+    """Second task has to wait for the first task due to insufficient resources"""
+    task_q, result_q = SpawnContext.Queue(), SpawnContext.Queue()
+    scheduler = MPITaskScheduler(task_q, result_q)
+
+    assert scheduler.available_nodes
+    assert len(scheduler.available_nodes) == 8
+
+    assert scheduler._free_node_counter.value == 8
+
+    mock_task_buffer = pack_res_spec_apply_message("func",
+                                                   "args",
+                                                   "kwargs",
+                                                   resource_specification={
+                                                       "num_nodes": 8,
+                                                       "ranks_per_node": 2
+                                                   })
+    task_package = {"task_id": 1, "buffer": mock_task_buffer}
+    scheduler.put_task(task_package)
+
+    assert scheduler._free_node_counter.value == 0
+    assert scheduler._backlog_queue.empty()
+
+    mock_task_buffer = pack_res_spec_apply_message("func",
+                                                   "args",
+                                                   "kwargs",
+                                                   resource_specification={
+                                                       "num_nodes": 8,
+                                                       "ranks_per_node": 2
+                                                   })
+    task_package = {"task_id": 2, "buffer": mock_task_buffer}
+    scheduler.put_task(task_package)
+
+    # Second task should now be in the backlog_queue
+    assert not scheduler._backlog_queue.empty()
+
+    # Confirm that the first task is available and has all 8 nodes provisioned
+    task_on_worker_side = task_q.get()
+    assert task_on_worker_side['task_id'] == 1
+    _, _, _, resource_spec = unpack_res_spec_apply_message(task_on_worker_side['buffer'])
+    assert len(resource_spec['MPI_NODELIST'].split(',')) == 8
+    assert task_q.empty()  # Confirm that task 2 is not yet scheduled
+
+    # Simulate worker returning result and the scheduler picking up result
+    result_pkl = pickle.dumps({"task_id": 1, "type": "result", "buffer": "RESULT BUF"})
+    result_q.put(result_pkl)
+    got_result = scheduler.get_result(True, 1)
+    assert got_result == result_pkl
+
+    # Now task2 must be scheduled
+    assert scheduler._backlog_queue.empty()
+
+    # Pop in a mock result
+    task_on_worker_side = task_q.get()
+    assert task_on_worker_side['task_id'] == 2
+    _, _, _, resource_spec = unpack_res_spec_apply_message(task_on_worker_side['buffer'])
+    assert len(resource_spec['MPI_NODELIST'].split(',')) == 8