/
base_plugin.py
71 lines (56 loc) · 2.14 KB
/
base_plugin.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
# -*- coding: utf-8 -*-
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import abc
from typing import Callable, Dict
from werkzeug import Response
class BasePlugin(abc.ABC):
"""Base plugin for cloud training tools endpoints.
The plugins support registering http handlers to be used for
AI Platform training jobs.
"""
@staticmethod
@abc.abstractmethod
def setup() -> None:
"""Run any setup code for the plugin before webserver is launched."""
raise NotImplementedError
@staticmethod
@abc.abstractmethod
def can_initialize() -> bool:
"""Check whether a plugin is able to be initialized.
Used for checking if correct dependencies are installed, system requirements, etc.
Returns:
Bool indicating whether the plugin can be initialized.
"""
raise NotImplementedError
@staticmethod
@abc.abstractmethod
def post_setup_check() -> bool:
"""Check if after initialization, we need to use the plugin.
Example: Web server only needs to run for main node for training, others
just need to have 'setup()' run to start the rpc server.
Returns:
A boolean indicating whether post setup checks pass.
"""
raise NotImplementedError
@abc.abstractmethod
def get_routes(self) -> Dict[str, Callable[..., Response]]:
"""Get the mapping from path to handler.
This is the method in which plugins can assign different routes to
different handlers.
Returns:
A mapping from a route to a handler.
"""
raise NotImplementedError