Skip to content

Commit 104c573

Browse files
krinsmankrinsman
committed
Minor refactor for legibility and maintainability
1 parent e44f2a3 commit 104c573

File tree

2 files changed

+121
-122
lines changed

2 files changed

+121
-122
lines changed

‎nbresuse/__init__.py

Lines changed: 1 addition & 122 deletions
Original file line numberDiff line numberDiff line change
@@ -1,40 +1,11 @@
11
import os
22

33
from tornado import ioloop
4-
from traitlets import Bool
5-
from traitlets import default
6-
from traitlets import Dict
7-
from traitlets import Float
8-
from traitlets import Int
9-
from traitlets import List
10-
from traitlets import TraitType
11-
from traitlets import Union
12-
from traitlets.config import Configurable
134

5+
from nbresuse.config import ResourceUseDisplay
146
from nbresuse.metrics import PSUtilMetricsLoader
157
from nbresuse.prometheus import PrometheusHandler
168

17-
try:
18-
# Traitlets >= 4.3.3
19-
from traitlets import Callable
20-
except ImportError:
21-
from .utils import Callable
22-
23-
24-
class PSUtilMetric(TraitType):
25-
"""A trait describing the format to specify a metric from the psutil package"""
26-
27-
info_text = "A dictionary specifying the function/method name, any keyword arguments, and if a named tuple is returned, which attribute of the named tuple to select"
28-
29-
def validate(self, obj, value):
30-
if isinstance(value, dict):
31-
keys = list(value.keys())
32-
if "name" in keys:
33-
keys.remove("name")
34-
if all(key in ["kwargs", "attribute"] for key in keys):
35-
return value
36-
self.error(obj, value)
37-
389

3910
def _jupyter_server_extension_paths():
4011
"""
@@ -57,98 +28,6 @@ def _jupyter_nbextension_paths():
5728
]
5829

5930

60-
class ResourceUseDisplay(Configurable):
61-
"""
62-
Holds server-side configuration for nbresuse
63-
"""
64-
65-
process_memory_metrics = List(
66-
trait=PSUtilMetric(),
67-
default_value=[{"name": "memory_info", "attribute": "rss"}],
68-
)
69-
70-
system_memory_metrics = List(
71-
trait=PSUtilMetric(),
72-
default_value=[{"name": "virtual_memory", "attribute": "total"}],
73-
)
74-
75-
process_cpu_metrics = List(
76-
trait=PSUtilMetric(),
77-
default_value=[{"name": "cpu_percent", "kwargs": {"interval": 0.05}}],
78-
)
79-
80-
system_cpu_metrics = List(
81-
trait=PSUtilMetric(), default_value=[{"name": "cpu_count"}]
82-
)
83-
84-
mem_warning_threshold = Float(
85-
default_value=0.1,
86-
help="""
87-
Warn user with flashing lights when memory usage is within this fraction
88-
memory limit.
89-
90-
For example, if memory limit is 128MB, `mem_warning_threshold` is 0.1,
91-
we will start warning the user when they use (128 - (128 * 0.1)) MB.
92-
93-
Set to 0 to disable warning.
94-
""",
95-
).tag(config=True)
96-
97-
mem_limit = Union(
98-
trait_types=[Int(), Callable()],
99-
help="""
100-
Memory limit to display to the user, in bytes.
101-
Can also be a function which calculates the memory limit.
102-
103-
Note that this does not actually limit the user's memory usage!
104-
105-
Defaults to reading from the `MEM_LIMIT` environment variable. If
106-
set to 0, the max memory available is displayed.
107-
""",
108-
).tag(config=True)
109-
110-
@default("mem_limit")
111-
def _mem_limit_default(self):
112-
return int(os.environ.get("MEM_LIMIT", 0))
113-
114-
track_cpu_percent = Bool(
115-
default_value=False,
116-
help="""
117-
Set to True in order to enable reporting of CPU usage statistics.
118-
""",
119-
).tag(config=True)
120-
121-
cpu_warning_threshold = Float(
122-
default_value=0.1,
123-
help="""
124-
Warn user with flashing lights when CPU usage is within this fraction
125-
CPU usage limit.
126-
127-
For example, if CPU limit is 150%, `cpu_warning_threshold` is 0.1,
128-
we will start warning the user when they use (150 - (150 * 0.1)) %.
129-
130-
Set to 0 to disable warning.
131-
""",
132-
).tag(config=True)
133-
134-
cpu_limit = Union(
135-
trait_types=[Float(), Callable()],
136-
default_value=0,
137-
help="""
138-
CPU usage limit to display to the user.
139-
140-
Note that this does not actually limit the user's CPU usage!
141-
142-
Defaults to reading from the `CPU_LIMIT` environment variable. If
143-
set to 0, the total CPU count available is displayed.
144-
""",
145-
).tag(config=True)
146-
147-
@default("cpu_limit")
148-
def _cpu_limit_default(self):
149-
return float(os.environ.get("CPU_LIMIT", 0))
150-
151-
15231
def load_jupyter_server_extension(nbapp):
15332
"""
15433
Called during notebook start

‎nbresuse/config.py

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
from traitlets import Bool
2+
from traitlets import default
3+
from traitlets import Dict
4+
from traitlets import Float
5+
from traitlets import Int
6+
from traitlets import List
7+
from traitlets import TraitType
8+
from traitlets import Union
9+
from traitlets.config import Configurable
10+
11+
try:
12+
# Traitlets >= 4.3.3
13+
from traitlets import Callable
14+
except ImportError:
15+
from .utils import Callable
16+
17+
class PSUtilMetric(TraitType):
18+
"""A trait describing the format to specify a metric from the psutil package"""
19+
20+
info_text = "A dictionary specifying the function/method name, any keyword arguments, and if a named tuple is returned, which attribute of the named tuple to select"
21+
22+
def validate(self, obj, value):
23+
if isinstance(value, dict):
24+
keys = list(value.keys())
25+
if "name" in keys:
26+
keys.remove("name")
27+
if all(key in ["kwargs", "attribute"] for key in keys):
28+
return value
29+
self.error(obj, value)
30+
31+
class ResourceUseDisplay(Configurable):
32+
"""
33+
Holds server-side configuration for nbresuse
34+
"""
35+
36+
process_memory_metrics = List(
37+
trait=PSUtilMetric(),
38+
default_value=[{"name": "memory_info", "attribute": "rss"}],
39+
)
40+
41+
system_memory_metrics = List(
42+
trait=PSUtilMetric(),
43+
default_value=[{"name": "virtual_memory", "attribute": "total"}],
44+
)
45+
46+
process_cpu_metrics = List(
47+
trait=PSUtilMetric(),
48+
default_value=[{"name": "cpu_percent", "kwargs": {"interval": 0.05}}],
49+
)
50+
51+
system_cpu_metrics = List(
52+
trait=PSUtilMetric(), default_value=[{"name": "cpu_count"}]
53+
)
54+
55+
mem_warning_threshold = Float(
56+
default_value=0.1,
57+
help="""
58+
Warn user with flashing lights when memory usage is within this fraction
59+
memory limit.
60+
61+
For example, if memory limit is 128MB, `mem_warning_threshold` is 0.1,
62+
we will start warning the user when they use (128 - (128 * 0.1)) MB.
63+
64+
Set to 0 to disable warning.
65+
""",
66+
).tag(config=True)
67+
68+
mem_limit = Union(
69+
trait_types=[Int(), Callable()],
70+
help="""
71+
Memory limit to display to the user, in bytes.
72+
Can also be a function which calculates the memory limit.
73+
74+
Note that this does not actually limit the user's memory usage!
75+
76+
Defaults to reading from the `MEM_LIMIT` environment variable. If
77+
set to 0, the max memory available is displayed.
78+
""",
79+
).tag(config=True)
80+
81+
@default("mem_limit")
82+
def _mem_limit_default(self):
83+
return int(os.environ.get("MEM_LIMIT", 0))
84+
85+
track_cpu_percent = Bool(
86+
default_value=False,
87+
help="""
88+
Set to True in order to enable reporting of CPU usage statistics.
89+
""",
90+
).tag(config=True)
91+
92+
cpu_warning_threshold = Float(
93+
default_value=0.1,
94+
help="""
95+
Warn user with flashing lights when CPU usage is within this fraction
96+
CPU usage limit.
97+
98+
For example, if CPU limit is 150%, `cpu_warning_threshold` is 0.1,
99+
we will start warning the user when they use (150 - (150 * 0.1)) %.
100+
101+
Set to 0 to disable warning.
102+
""",
103+
).tag(config=True)
104+
105+
cpu_limit = Union(
106+
trait_types=[Float(), Callable()],
107+
default_value=0,
108+
help="""
109+
CPU usage limit to display to the user.
110+
111+
Note that this does not actually limit the user's CPU usage!
112+
113+
Defaults to reading from the `CPU_LIMIT` environment variable. If
114+
set to 0, the total CPU count available is displayed.
115+
""",
116+
).tag(config=True)
117+
118+
@default("cpu_limit")
119+
def _cpu_limit_default(self):
120+
return float(os.environ.get("CPU_LIMIT", 0))

0 commit comments

Comments
 (0)