feat: support set metadir as a fold from the cloud

Author: pwwang

.gitignore

@@ -158,4 +158,6 @@ docs/api/
 test.py
 
 /local_tests/
-t.ipynb
+t.ipynb
+
+gac_key.json

poetry.lock

@@ -8,22 +8,23 @@ homepage = "https://github.com/pwwang/xqute"
 repository  = "https://github.com/pwwang/xqute"
 readme = "README.md"
 
+[tool.poetry.build]
+generate-setup-file = true
+
 [tool.poetry.dependencies]
-python = "^3.8"
+python = "^3.9"
 simplug = "^0.4"
 diot = "^0.2"
 rich = "^13"
 uvloop = "^0"
-aiopath = [
-    {version = "0.5.*", python = "<3.10"},
-    {version = "0.6.*", python = ">=3.10,<3.12"},
-    {version = "0.7.*", python = ">=3.12"},
-]
+cloudpathlib = "^0.20"
+google-cloud-storage = "^2.19.0"
+argx = "^0.2.11"
 
-[tool.poetry.build]
-generate-setup-file = true
+[tool.poetry.extras]
+meta-gs = ["google-cloud-storage"]
 
-[tool.poetry.dev-dependencies]
+[tool.poetry.group.dev.dependencies]
 pytest = "^8"
 pytest-cov = "^5"
 pytest-asyncio = "^0"

pyproject.toml

@@ -0,0 +1,3 @@
+import os
+
+os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "./gac_key.json"

tests/conftest.py

@@ -1,5 +1,7 @@
+import asyncio
 import pytest
 
+from xqute.defaults import JobStatus
 from xqute.schedulers.local_scheduler import LocalJob, LocalScheduler
 
 
@@ -12,11 +14,26 @@ async def test_scheduler():
     assert isinstance(pid, int)
 
 
+@pytest.mark.asyncio
+async def test_scheduler_with_meta_gs():
+    job = LocalJob(
+        0,
+        ["echo", 1],
+        metadir="gs://handy-buffer-287000.appspot.com/xqute_local_test",
+    )
+
+    scheduler = LocalScheduler(1)
+    pid = await scheduler.submit_job(job)
+
+    assert isinstance(pid, int)
+    assert job.stdout_file.read_text() == "1\n"
+
+
 @pytest.mark.asyncio
 async def test_immediate_submission_failure():
 
     class BadLocalJob(LocalJob):
-        async def wrapped_script(self, scheduler):
+        def wrapped_script(self, scheduler):
             wrapt_script = self.metadir / f"job.wrapped.{scheduler.name}"
             wrapt_script.write_text("sleep 1; bad_non_existent_command")
             return wrapt_script
@@ -30,3 +47,21 @@ async def wrapped_script(self, scheduler):
         RuntimeError, match="bad_non_existent_command: command not found"
     ):
         await scheduler.submit_job(job)
+
+
+@pytest.mark.asyncio
+async def test_killing_running_jobs(caplog):
+    job1 = LocalJob(0, ["sleep", "10"])
+    job2 = LocalJob(1, ["sleep", "10"])
+
+    scheduler = LocalScheduler(2)
+    await scheduler.submit_job_and_update_status(job1)
+    await scheduler.submit_job_and_update_status(job2)
+
+    await asyncio.sleep(1)
+    await scheduler.kill_running_jobs([job1, job2])
+
+    assert job1.status == JobStatus.FINISHED
+    assert job2.status == JobStatus.FINISHED
+    assert job1.rc != 0
+    assert job2.rc != 0

tests/test_local_scheduler.py

@@ -0,0 +1,193 @@
+import asyncio
+import pytest
+import subprocess
+from argx import Namespace
+from unittest.mock import patch, MagicMock, AsyncMock
+from xqute import __version__
+from xqute.__main__ import (
+    _update_status,
+    check_version_and_plugins,
+    launch_job,
+    main,
+    SchedulerXquteInconsistencyWarning,
+)
+from xqute.defaults import JobStatus
+
+# FILE: xqute/test___main__.py
+
+
+def test_update_status(tmp_path):
+    metadir = tmp_path / "metadir"
+    metadir.mkdir()
+    status = JobStatus.RUNNING
+    _update_status(str(metadir), status)
+    assert (metadir / "job.status").read_text() == str(status)
+
+
+def test_check_version_and_plugins():
+    args = MagicMock()
+    args.version = "1.0.0"
+    args.plugin = ["plugin1:1.0.0", "plugin2:2.0.0"]
+    plugin1 = Namespace(name="plugin1", version="1.0.0")
+    plugin2 = Namespace(name="plugin2", version="2.0.0")
+
+    with patch("xqute.__main__.__version__", "1.0.0"):
+        with patch(
+            "xqute.__main__.plugin.get_enabled_plugins",
+            return_value={"plugin1": "1.0.0", "plugin2": "2.0.0"},
+        ):
+            check_version_and_plugins(args)
+
+    with patch("xqute.__main__.__version__", "1.0.0"):
+        with patch(
+            "xqute.__main__.plugin.get_enabled_plugins",
+            return_value={"plugin1": "1.0.0"},
+        ):
+            with pytest.warns(SchedulerXquteInconsistencyWarning):
+                check_version_and_plugins(args)
+
+    with patch("xqute.__main__.__version__", "1.0.0"):
+        with patch(
+            "xqute.__main__.plugin.get_enabled_plugins",
+            return_value={"plugin1": plugin1, "plugin2": plugin2},
+        ):
+            with pytest.warns(SchedulerXquteInconsistencyWarning):
+                check_version_and_plugins(args)
+
+
+def test_check_version_and_plugins_version_mismatch():
+    args = MagicMock()
+    args.version = "1.0.0"
+
+    with patch("xqute.__main__.__version__", "2.0.0"):
+        with pytest.warns(
+            SchedulerXquteInconsistencyWarning,
+            match="The xqute on the scheduler system is 2.0.0",
+        ):
+            check_version_and_plugins(args)
+
+    args.plugin = ["plugin1:1.0.0"]
+    with patch("xqute.__main__.__version__", "1.0.0"):
+        with patch(
+            "xqute.__main__.plugin.get_enabled_plugins",
+            return_value={"plugin1": "2.0.0"},
+        ):
+            with pytest.warns(
+                SchedulerXquteInconsistencyWarning,
+                match="Xqute plugin version mismatch: plugin1 1.0.0 != 2.0.0",
+            ):
+                check_version_and_plugins(args)
+
+
+@pytest.mark.asyncio
+async def test_launch_job(tmp_path):
+    args = MagicMock()
+    args.metadir = str(tmp_path / "metadir")
+    args.cmd = ["echo", "hello"]
+    (tmp_path / "metadir").mkdir()
+
+    with patch(
+        "xqute.__main__.plugin.hooks.on_jobsched_started", new_callable=AsyncMock
+    ):
+        with patch(
+            "xqute.__main__.plugin.hooks.on_jobsched_ended", new_callable=AsyncMock
+        ):
+            await launch_job(args)
+
+    assert (tmp_path / "metadir" / "job.stdout").read_text() == "hello\n"
+    assert (tmp_path / "metadir" / "job.rc").read_text() == "0"
+    assert (tmp_path / "metadir" / "job.status").read_text() == str(JobStatus.FINISHED)
+
+
+@pytest.mark.asyncio
+async def test_launch_job_failed(tmp_path):
+    metadir = tmp_path / "metadir"
+    metadir.mkdir()
+    args = Namespace(
+        metadir=str(metadir),
+        cmd=["bash", "-c", "exit 1"],
+        scheduler="local",
+        version=__version__,
+        plugin=[],
+    )
+
+    with pytest.raises(SystemExit):
+        await launch_job(args)
+    assert (metadir / "job.rc").read_text() != "0"
+
+    args = Namespace(
+        metadir=str(metadir),
+        cmd=["__NoNeXisT__"],
+        scheduler="local",
+        version=__version__,
+        plugin=[],
+    )
+
+    with pytest.raises(SystemExit):
+        await launch_job(args)
+    assert (metadir / "job.rc").read_text() != "0"
+
+
+@pytest.mark.asyncio
+async def test_launch_job_cancelled(tmp_path):
+    args = MagicMock()
+    args.metadir = str(tmp_path / "metadir")
+    args.cmd = ["echo", "hello"]
+    (tmp_path / "metadir").mkdir()
+
+    with patch(
+        "xqute.__main__.plugin.hooks.on_jobsched_started", new_callable=AsyncMock
+    ):
+        with patch(
+            "xqute.__main__.plugin.hooks.on_jobsched_ended", new_callable=AsyncMock
+        ):
+            with patch(
+                "asyncio.create_subprocess_exec", side_effect=asyncio.CancelledError
+            ):
+                with pytest.raises(asyncio.CancelledError):
+                    await launch_job(args)
+
+    assert (tmp_path / "metadir" / "job.rc").read_text() == "1"
+    assert (tmp_path / "metadir" / "job.status").read_text() == str(JobStatus.FAILED)
+
+
+@pytest.mark.asyncio
+async def test_main(tmp_path):
+    metadir = tmp_path / "metadir"
+    metadir.mkdir()
+    args = Namespace(
+        metadir=str(metadir),
+        cmd=["echo", "hello"],
+        scheduler="local",
+        version=__version__,
+        plugin=[],
+    )
+
+    await main(args)
+    assert (metadir / "job.stdout").read_text() == "hello\n"
+
+
+@pytest.mark.asyncio
+async def test_cli(tmp_path):
+    metadir = tmp_path / "metadir"
+    metadir.mkdir()
+
+    # Use the real command line arguments
+    p = subprocess.run(
+        [
+            "python",
+            "-m",
+            "xqute",
+            "++scheduler",
+            "local",
+            "++version",
+            __version__,
+            "++metadir",
+            str(metadir),
+            "++cmd",
+            "echo",
+            "hello",
+        ]
+    )
+    assert p.returncode == 0
+    assert (metadir / "job.stdout").read_text() == "hello\n"

tests/test_main.py

@@ -9,14 +9,28 @@
 MOCKS = Path(__file__).parent / "mocks"
 
 
-def setup_module():
-    qsub = str(MOCKS / "qsub")
-    qdel = str(MOCKS / "qdel")
-    qstat = str(MOCKS / "qstat")
+@pytest.fixture
+def qsub():
+    cmd = str(MOCKS / "qsub")
+    st = os.stat(cmd)
+    os.chmod(cmd, st.st_mode | stat.S_IEXEC)
+    return cmd
 
-    for qcmd in (qsub, qdel, qstat):
-        st = os.stat(str(qcmd))
-        os.chmod(str(qcmd), st.st_mode | stat.S_IEXEC)
+
+@pytest.fixture
+def qdel():
+    cmd = str(MOCKS / "qdel")
+    st = os.stat(cmd)
+    os.chmod(cmd, st.st_mode | stat.S_IEXEC)
+    return cmd
+
+
+@pytest.fixture
+def qstat():
+    cmd = str(MOCKS / "qstat")
+    st = os.stat(cmd)
+    os.chmod(cmd, st.st_mode | stat.S_IEXEC)
+    return cmd
 
 
 @pytest.mark.asyncio
@@ -26,23 +40,20 @@ async def test_job():
         forks=1, qsub_notify=True, sge_l=["vmem=2G", "gpu=1"], sge_m="abe"
     )
     assert (
-        await job.wrapped_script(scheduler)
-        == DEFAULT_JOB_METADIR / "0" / "job.wrapped.sge"
+        job.wrapped_script(scheduler)
+        == Path(DEFAULT_JOB_METADIR) / "0" / "job.wrapped.sge"
     )
 
-    shebang = job.shebang(scheduler)
-    assert "#$ -notify" in shebang
-    assert "#$ -l vmem=2G" in shebang
-    assert "#$ -l gpu=1" in shebang
-    assert "#$ -m abe" in shebang
+    script = job.wrap_script(scheduler)
+    assert "#$ -notify" in script
+    assert "#$ -l vmem=2G" in script
+    assert "#$ -l gpu=1" in script
+    assert "#$ -m abe" in script
 
 
 @pytest.mark.asyncio
-async def test_scheduler(capsys):
+async def test_scheduler(capsys, qsub, qdel, qstat):
     job = SgeJob(0, ["echo", 1])
-    qsub = str(MOCKS / "qsub")
-    qdel = str(MOCKS / "qdel")
-    qstat = str(MOCKS / "qstat")
 
     scheduler = SgeScheduler(1, qsub=qsub, qdel=qdel, qstat=qstat)
     assert await scheduler.submit_job(job) == "613815"
@@ -61,13 +72,10 @@ async def test_scheduler(capsys):
 
 
 @pytest.mark.asyncio
-async def test_submission_failure(capsys):
+async def test_submission_failure(capsys, qdel, qstat):
     job = SgeJob(0, ["echo", 1])
-    qsub = str(MOCKS / "no_such_qsub")
-    qdel = str(MOCKS / "qdel")
-    qstat = str(MOCKS / "qstat")
 
-    scheduler = SgeScheduler(1, qsub=qsub, qdel=qdel, qstat=qstat)
+    scheduler = SgeScheduler(1, qsub="no_such_qsub", qdel=qdel, qstat=qstat)
 
     assert await scheduler.submit_job_and_update_status(job) is None
     assert await scheduler.job_is_running(job) is False