Code review changes

Author: ChrisPaulBennett

cylc/flow/remote.py

@@ -34,6 +34,7 @@
 )
 import sys
 from typing import (
+    TYPE_CHECKING,
     Any,
     Dict,
     List,
@@ -56,6 +57,9 @@
 )
 from cylc.flow.util import format_cmd
 
+if TYPE_CHECKING:
+    import psutil
+
 
 def get_proc_ancestors():
     """Return list of parent PIDs back to init."""
@@ -76,7 +80,10 @@ def get_proc_ancestors():
         pid = ppid
 
 
-async def watch_and_kill(proc, interval=None):
+async def watch_and_kill(
+    proc: 'psutil.Process',
+    interval: float | None = None,
+):
     """Watch a process and kill it if any of its parent processes change.
 
     Processes exist in a tree which inherits from the process with PID 1.

cylc/flow/scripts/cat_log.py

@@ -61,6 +61,7 @@
   $ cylc cat-log foo//2020/bar -m f
 """
 
+import asyncio
 import os
 from contextlib import suppress
 from glob import glob
@@ -70,10 +71,12 @@
 import sys
 from typing import TYPE_CHECKING
 
+from psutil import Process
+
 from cylc.flow.exceptions import InputError
 import cylc.flow.flags
 from cylc.flow.hostuserutil import is_remote_platform
-from cylc.flow.id_cli import parse_id
+from cylc.flow.id_cli import parse_id_async
 from cylc.flow.log_level import verbosity_to_opts
 from cylc.flow.option_parsers import (
     ID_MULTI_ARG_DOC,
@@ -244,7 +247,7 @@ def _check_fs_path(path):
         )
 
 
-def view_log(
+async def view_log(
     logpath,
     mode,
     tailer_tmpl,
@@ -308,7 +311,7 @@ def view_log(
         proc = Popen(shlex.split(cmd), stdin=DEVNULL)  # nosec
         # * batchview command is user configurable
         with suppress(KeyboardInterrupt):
-            watch_and_kill(proc)
+            await watch_and_kill(Process(proc.pid))
         return proc.wait()
 
 
@@ -414,10 +417,10 @@ def main(
 ):
     """Wrapper around the main script for simpler testing.
     """
-    _main(parser, options, *ids, color=color)
+    asyncio.run(_main(parser, options, *ids, color=color))
 
 
-def _main(
+async def _main(
     parser: COP,
     options: 'Values',
     *ids,
@@ -445,7 +448,7 @@ def _main(
             batchview_cmd = options.remote_args[3]
         except IndexError:
             batchview_cmd = None
-        res = view_log(
+        res = await view_log(
             logpath,
             mode,
             tail_tmpl,
@@ -458,7 +461,7 @@ def _main(
             sys.exit(res)
         return
 
-    workflow_id, tokens, _ = parse_id(*ids, constraint='mixed')
+    workflow_id, tokens, _ = await parse_id_async(*ids, constraint='mixed')
 
     # Get long-format mode.
     try:
@@ -522,7 +525,7 @@ def _main(
         tail_tmpl = os.path.expandvars(
             get_platform()["tail command template"]
         )
-        out = view_log(
+        out = await view_log(
             log_file_path,
             mode,
             tail_tmpl,
@@ -679,7 +682,7 @@ def _main(
             # Local task job or local job log.
             logpath = os.path.join(local_log_dir, options.filename)
             tail_tmpl = os.path.expandvars(platform["tail command template"])
-            out = view_log(
+            out = await view_log(
                 logpath,
                 mode,
                 tail_tmpl,

cylc/flow/scripts/cylc.py

@@ -653,7 +653,7 @@ def main() -> None:  # pragma: no cover
 
 
 def _main(opts, cmd_args) -> int:
-    """Implemnent the Cylc CLI.
+    """Implement the Cylc CLI.
 
     Returns the exit code as an integer.
     """

cylc/flow/scripts/message.py

@@ -93,12 +93,13 @@
 """
 
 
+import asyncio
 from logging import getLevelName, INFO
 import os
 import sys
 from typing import TYPE_CHECKING
 
-from cylc.flow.id_cli import parse_id
+from cylc.flow.id_cli import parse_id_async
 from cylc.flow.option_parsers import (
     WORKFLOW_ID_ARG_DOC,
     CylcOptionParser as COP
@@ -142,6 +143,10 @@ def main(parser: COP, options: 'Values', *args: str) -> None:
     if not args:
         parser.error('No message supplied')
         return
+    asyncio.run(_main(options, args))
+
+
+async def _main(options, args):
     if len(args) <= 2:
         # BACK COMPAT: args <= 2
         # from:
@@ -160,7 +165,7 @@ def main(parser: COP, options: 'Values', *args: str) -> None:
         message_strs = list(args)
     else:
         workflow_id, job_id, *message_strs = args
-        workflow_id, *_ = parse_id(
+        workflow_id, *_ = await parse_id_async(
             workflow_id,
             constraint='workflows',
         )
@@ -198,4 +203,4 @@ def main(parser: COP, options: 'Values', *args: str) -> None:
             messages.append([options.severity, message_str.strip()])
         else:
             messages.append([getLevelName(INFO), message_str.strip()])
-    record_messages(workflow_id, job_id, messages)
+    await record_messages(workflow_id, job_id, messages)

cylc/flow/scripts/profiler.py

@@ -21,13 +21,12 @@
 """
 
 import asyncio
+from contextlib import suppress
 import json
 import os
 import re
-import sys
 import signal
 import psutil
-import functools
 
 from pathlib import Path
 from dataclasses import dataclass
@@ -53,20 +52,25 @@ class Process:
     cgroup_version: int
 
 
-def stop_profiler(process, comms_timeout, *_args):
+async def stop_profiler(process, comms_timeout, tasks, *_args):
     """Stop the profiler and return its data to the scheduler.
 
     This function will be executed when the profiler receives a stop signal.
     """
+    # stop the profiler
+    for task in tasks:
+        task.cancel()
+
+    # extract the stats
     profiler_data = get_profiler_data(process)
 
-    record_messages(
+    # send a task message to the scheduler / write message to job.status file
+    await record_messages(
         os.environ['CYLC_WORKFLOW_ID'],
         os.environ['CYLC_TASK_JOB'],
         [['DEBUG', f'_cylc_profiler: {json.dumps(profiler_data)}']],
         comms_timeout=comms_timeout,
     )
-    sys.exit(0)
 
 
 def get_profiler_data(process):
@@ -219,7 +223,7 @@ async def profile(_process: Process, delay, keep_looping=lambda: True):
     while keep_looping():
         # Polling the cgroup for memory and keeping track of the max rss value
         max_rss = parse_memory_file(_process)
-        if max_rss > _process.max_rss:
+        if max_rss is not None and max_rss > _process.max_rss:
             _process.max_rss = max_rss
         await asyncio.sleep(delay)
 
@@ -244,29 +248,36 @@ def get_option_parser() -> COP:
 @cli_function(get_option_parser)
 def main(_parser: COP, options) -> None:
     """CLI main."""
-    asyncio.run(_main(options))
+    with suppress(SystemExit, asyncio.exceptions.CancelledError, Exception):
+        asyncio.run(_main(options))
 
 
 async def _main(options) -> None:
     # get cgroup information
     process = get_cgroup_paths(options.cgroup_location)
 
+    # list of asyncio tasks
+    tasks = []
+
     # Register the stop_profiler function with the signal library
     # The signal library doesn't work with asyncio, so we have to use the
     # loop's add_signal_handler function instead
     loop = asyncio.get_running_loop()
     for sig in (signal.SIGINT, signal.SIGHUP, signal.SIGTERM):
         loop.add_signal_handler(
             sig,
-            functools.partial(stop_profiler, process, options.comms_timeout)
+            lambda: asyncio.create_task(
+                stop_profiler(process, options.comms_timeout, tasks)
+            ),
         )
 
     # the profiler will run until one of these coroutines calls `sys.exit`:
-    await asyncio.gather(
+    tasks.extend([
         # run the profiler itself
-        profile(process, options.delay),
+        asyncio.create_task(profile(process, options.delay)),
 
         # kill the profiler if its PPID changes
         # (i.e, if the job exits before the profiler does)
-        watch_and_kill(psutil.Process(os.getpid())),
-    )
+        asyncio.create_task(watch_and_kill(psutil.Process(os.getpid()))),
+    ])
+    await asyncio.gather(*tasks)

cylc/flow/task_message.py

@@ -92,7 +92,7 @@ def split_run_signal(message: str) -> tuple[str, str | None]:
     return prefix, signal[0] if signal else None
 
 
-def record_messages(
+async def record_messages(
     workflow: str,
     job_id: str,
     messages: List[list],
@@ -115,7 +115,8 @@ def record_messages(
         override_use_utc=(os.getenv('CYLC_UTC') == 'True'))
     write_messages(workflow, job_id, messages, event_time)
     if get_comms_method() != CommsMeth.POLL:
-        send_messages(workflow, job_id, messages, event_time, comms_timeout)
+        await send_messages(workflow, job_id, messages,
+                            event_time, comms_timeout)
 
 
 def write_messages(workflow, job_id, messages, event_time):
@@ -131,7 +132,7 @@ def write_messages(workflow, job_id, messages, event_time):
     _append_job_status_file(workflow, job_id, event_time, messages)
 
 
-def send_messages(
+async def send_messages(
     workflow: str,
     job_id: str,
     messages: List[list],
@@ -163,7 +164,7 @@ def send_messages(
             'messages': messages,
         }
     }
-    pclient('graphql', mutation_kwargs)
+    await pclient.async_request('graphql', mutation_kwargs)
 
 
 def _append_job_status_file(workflow, job_id, event_time, messages):

tests/functional/cylc-cat-log/12-delete-kill.t

@@ -19,7 +19,7 @@
 # or when tail is killed.
 
 . "$(dirname "$0")/test_header"
-set_test_number 2
+set_test_number 1
 
 # Get PID of tail cmd given the parent cat-log PPID
 get_tail_pid() {
@@ -33,6 +33,8 @@ __EOF__
 log_file="${WORKFLOW_RUN_DIR}/log/foo.log"
 echo "Hello, Mr. Thompson" > "$log_file"
 
+export CYLC_PROC_POLL_INTERVAL=0.5
+
 TEST_NAME="${TEST_NAME_BASE}-delete"
 cylc cat-log --mode=tail "$WORKFLOW_NAME" -f foo.log 2>&1 &
 cat_log_pid="$!"

tests/integration/scripts/test_cat_log.py

@@ -39,16 +39,18 @@ def brokendir(run_dir):
     shutil.rmtree(brokendir)
 
 
-def test_fail_no_file(flow):
+@pytest.mark.asyncio
+async def test_fail_no_file(flow):
     """It produces a helpful error if there is no workflow log file.
     """
     parser = cat_log_gop()
     id_ = flow({})
     with pytest.raises(InputError, match='Log file not found.'):
-        cat_log(parser, Options(parser)(), id_)
+        await cat_log(parser, Options(parser)(), id_)
 
 
-def test_fail_rotation_out_of_range(flow):
+@pytest.mark.asyncio
+async def test_fail_rotation_out_of_range(flow):
     """It produces a helpful error if rotation number > number of log files.
     """
     parser = cat_log_gop()
@@ -60,15 +62,16 @@ def test_fail_rotation_out_of_range(flow):
     (logpath / '01-start-01.log').touch()
 
     with pytest.raises(SystemExit):
-        cat_log(parser, Options(parser)(rotation_num=0), id_)
+        await cat_log(parser, Options(parser)(rotation_num=0), id_)
 
     msg = r'--rotation 1 invalid \(max value is 0\)'
 
     with pytest.raises(InputError, match=msg):
-        cat_log(parser, Options(parser)(rotation_num=1), id_)
+        await cat_log(parser, Options(parser)(rotation_num=1), id_)
 
 
-def test_bad_workflow(run_dir):
+@pytest.mark.asyncio
+async def test_bad_workflow(run_dir):
     """Test "cylc cat-log" with bad workflow name."""
     parser = cat_log_gop()
     msg = re.compile(
@@ -77,15 +80,16 @@ def test_bad_workflow(run_dir):
         re.MULTILINE
     )
     with pytest.raises(InputError, match=msg):
-        cat_log(parser, Options(parser)(filename='l'), BAD_NAME)
+        await cat_log(parser, Options(parser)(filename='l'), BAD_NAME)
 
 
-def test_bad_workflow2(run_dir, brokendir, capsys):
+@pytest.mark.asyncio
+async def test_bad_workflow2(run_dir, brokendir, capsys):
     """Check a non existent file in a valid workflow results in error.
     """
     parser = cat_log_gop()
     with pytest.raises(SystemExit, match='1'):
-        cat_log(
+        await cat_log(
             parser,
             Options(parser)(filename='j'),
             BAD_NAME
@@ -96,12 +100,13 @@ def test_bad_workflow2(run_dir, brokendir, capsys):
     assert capsys.readouterr().err == msg
 
 
-def test_bad_task_dir(run_dir, brokendir, capsys):
+@pytest.mark.asyncio
+async def test_bad_task_dir(run_dir, brokendir, capsys):
     """Check a non-existent job log dir in a valid workflow results in error.
     """
     parser = cat_log_gop()
     with pytest.raises(SystemExit, match='1'):
-        cat_log(
+        await cat_log(
             parser,
             Options(parser)(mode='list-dir'),
             BAD_NAME + "//1/foo"