-
Notifications
You must be signed in to change notification settings - Fork 4
/
pre_run.py
92 lines (70 loc) · 3.19 KB
/
pre_run.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import os
from _py_abc import ABCMeta
from abc import abstractmethod
from pypads import logger
from pypads.functions.loggers.base_logger import FunctionHolder
from pypads.functions.loggers.mixins import DependencyMixin, OrderMixin, TimedCallableMixin, IntermediateCallableMixin, \
ConfigurableCallableMixin, DefensiveCallableMixin
class PreRunFunction(DefensiveCallableMixin, IntermediateCallableMixin, FunctionHolder, TimedCallableMixin,
DependencyMixin, OrderMixin,
ConfigurableCallableMixin):
"""
This class should be used to define new pre run functions
"""
__metaclass__ = ABCMeta
@abstractmethod
def __init__(self, *args, fn=None, **kwargs):
super().__init__(*args, fn=fn, **kwargs)
if self._fn is None:
self._fn = self._call
@abstractmethod
def _call(self, pads, *args, **kwargs):
"""
Function where to add you custom code to execute before starting the run.
:param pads: the current instance of PyPads.
"""
return NotImplementedError()
def __real_call__(self, *args, **kwargs):
from pypads.pypads import get_current_pads
return super().__real_call__(get_current_pads(), *args, **kwargs)
def _handle_error(self, *args, ctx, _pypads_env, error, **kwargs):
logger.warning("Couldn't execute " + str(self) + ", because of exception: " + str(error))
class RunInfo(PreRunFunction):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
@staticmethod
def _needed_packages():
return ["pip"]
def _call(self, pads, *args, **kwargs):
logger.info("Tracking execution to run with id " + pads.api.active_run().info.run_id)
# Execute pip freeze
try:
# noinspection PyProtectedMember,PyPackageRequirements
from pip._internal.operations import freeze
except ImportError: # pip < 10.0
# noinspection PyUnresolvedReferences,PyPackageRequirements
from pip.operations import freeze
pads.api.log_mem_artifact("pip_freeze", "\n".join(freeze.freeze()))
class RunLogger(PreRunFunction):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def _call(self, pads, *args, **kwargs):
from pypads.base import PypadsApi
_api: PypadsApi = pads.api
from pypads.logging_util import get_temp_folder
folder = get_temp_folder()
# TODO loguru has problems with multiprocessing / make rotation configurable etc
from pypads.pads_loguru import logger_manager
lid = logger_manager.add(os.path.join(folder, "run_" + _api.active_run().info.run_id + ".log"),
rotation="50 MB",
enqueue=True)
import glob
def remove_logger(pads, *args, **kwargs):
try:
from pypads.pads_loguru import logger_manager
logger_manager.remove(lid)
except Exception:
pass
for file in glob.glob(os.path.join(folder, "run_*.log")):
pads.api.log_artifact(file)
_api.register_post_fn("logger_" + str(lid), remove_logger)