pytorch-lightning

1
# Copyright The Lightning AI team.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14

15
import os
16
import signal
17
import subprocess
18
import sys
19
from pathlib import Path
20
from typing import Any, Dict, List, Optional, Union
21

22
from lightning.app.core.work import LightningWork
23
from lightning.app.utilities.app_helpers import Logger, _collect_child_process_pids
24
from lightning.app.utilities.tracer import Tracer
25

26
logger = Logger(__name__)
27

28

29
class PopenPythonScript(LightningWork):
30
    def on_before_run(self):
31
        """Called before the python script is executed."""
32

33
    def on_after_run(self):
34
        """Called after the python script is executed."""
35

36
    def configure_tracer(self) -> Tracer:
37
        """Override this hook to customize your tracer when running PythonScript with ``mode=tracer``."""
38
        return Tracer()
39

40
    def __init__(
41
        self,
42
        script_path: Union[str, Path],
43
        script_args: Optional[Union[str, List[str]]] = None,
44
        env: Optional[Dict] = None,
45
        **kwargs: Any,
46
    ):
47
        """The PopenPythonScript component enables to easily run a python script within a subprocess.
48

49
        Arguments:
50
            script_path: Path of the python script to run.
51
            script_path: The arguments to be passed to the script.
52
            env: Environment variables to be passed to the script.
53
            kwargs: LightningWork keyword arguments.
54

55
        Raises:
56
            FileNotFoundError: If the provided `script_path` doesn't exists.
57

58
        Example:
59

60
            >>> from lightning.app.components.python import PopenPythonScript
61
            >>> f = open("a.py", "w")
62
            >>> f.write("print('Hello World !')")
63
            22
64
            >>> f.close()
65
            >>> python_script = PopenPythonScript("a.py")
66
            >>> python_script.run()
67
            >>> os.remove("a.py")
68

69
        In this example, the script will be launch with the :class:`~subprocess.Popen`.
70

71
        .. literalinclude:: ../../../../examples/app/components/python/component_popen.py
72
            :language: python
73

74
        """
75
        super().__init__(**kwargs)
76
        if not os.path.exists(script_path):
77
            raise FileNotFoundError(f"The provided `script_path` {script_path}` wasn't found.")
78
        self.script_path = str(script_path)
79
        if isinstance(script_args, str):
80
            script_args = script_args.split(" ")
81
        self.script_args = script_args if script_args else []
82
        self.env = env
83
        self.pid = None
84
        self.exit_code = None
85

86
    def run(self) -> None:
87
        self.on_before_run()
88
        self._run_with_subprocess_popen()
89
        self.on_after_run()
90
        return
91

92
    def _run_with_subprocess_popen(self) -> None:
93
        cmd = [sys.executable] + [self.script_path] + self.script_args
94

95
        with subprocess.Popen(
96
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=0, close_fds=True, env=self.env
97
        ) as proc:
98
            self.pid = proc.pid
99
            if proc.stdout:
100
                with proc.stdout:
101
                    for line in iter(proc.stdout.readline, b""):
102
                        logger.info("%s", line.decode().rstrip())
103

104
            self.exit_code = proc.wait()
105
            if self.exit_code != 0:
106
                raise Exception(self.exit_code)
107

108
    def on_exit(self):
109
        for child_pid in _collect_child_process_pids(os.getpid()):
110
            os.kill(child_pid, signal.SIGTERM)
111

112

113
__all__ = ["PopenPythonScript"]
114