digitalkin 0.2.11__py3-none-any.whl → 0.2.13__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- digitalkin/__version__.py +1 -1
- digitalkin/grpc_servers/_base_server.py +15 -17
- digitalkin/grpc_servers/module_server.py +9 -10
- digitalkin/grpc_servers/module_servicer.py +108 -85
- digitalkin/grpc_servers/registry_server.py +3 -6
- digitalkin/grpc_servers/registry_servicer.py +18 -19
- digitalkin/grpc_servers/utils/grpc_client_wrapper.py +3 -5
- digitalkin/logger.py +45 -1
- digitalkin/models/module/module.py +1 -0
- digitalkin/modules/_base_module.py +47 -6
- digitalkin/modules/job_manager/base_job_manager.py +139 -0
- digitalkin/modules/job_manager/job_manager_models.py +44 -0
- digitalkin/modules/job_manager/single_job_manager.py +218 -0
- digitalkin/modules/job_manager/taskiq_broker.py +173 -0
- digitalkin/modules/job_manager/taskiq_job_manager.py +213 -0
- digitalkin/services/base_strategy.py +3 -1
- digitalkin/services/cost/cost_strategy.py +64 -16
- digitalkin/services/cost/default_cost.py +95 -12
- digitalkin/services/cost/grpc_cost.py +149 -60
- digitalkin/services/filesystem/default_filesystem.py +5 -6
- digitalkin/services/filesystem/filesystem_strategy.py +3 -2
- digitalkin/services/filesystem/grpc_filesystem.py +31 -26
- digitalkin/services/services_config.py +6 -5
- digitalkin/services/setup/__init__.py +1 -0
- digitalkin/services/setup/default_setup.py +10 -12
- digitalkin/services/setup/grpc_setup.py +8 -10
- digitalkin/services/storage/default_storage.py +13 -6
- digitalkin/services/storage/grpc_storage.py +25 -9
- digitalkin/services/storage/storage_strategy.py +3 -2
- digitalkin/utils/arg_parser.py +5 -48
- digitalkin/utils/development_mode_action.py +51 -0
- {digitalkin-0.2.11.dist-info → digitalkin-0.2.13.dist-info}/METADATA +43 -12
- {digitalkin-0.2.11.dist-info → digitalkin-0.2.13.dist-info}/RECORD +40 -33
- {digitalkin-0.2.11.dist-info → digitalkin-0.2.13.dist-info}/WHEEL +1 -1
- modules/cpu_intensive_module.py +271 -0
- modules/minimal_llm_module.py +200 -56
- modules/storage_module.py +5 -6
- modules/text_transform_module.py +1 -1
- digitalkin/modules/job_manager.py +0 -176
- {digitalkin-0.2.11.dist-info → digitalkin-0.2.13.dist-info}/licenses/LICENSE +0 -0
- {digitalkin-0.2.11.dist-info → digitalkin-0.2.13.dist-info}/top_level.txt +0 -0
|
@@ -1,8 +1,8 @@
|
|
|
1
1
|
"""This module contains the abstract base class for cost strategies."""
|
|
2
2
|
|
|
3
3
|
from abc import ABC, abstractmethod
|
|
4
|
-
from enum import Enum
|
|
5
|
-
from typing import
|
|
4
|
+
from enum import Enum
|
|
5
|
+
from typing import Literal
|
|
6
6
|
|
|
7
7
|
from pydantic import BaseModel
|
|
8
8
|
|
|
@@ -10,36 +10,84 @@ from digitalkin.services.base_strategy import BaseStrategy
|
|
|
10
10
|
|
|
11
11
|
|
|
12
12
|
class CostType(Enum):
|
|
13
|
-
"""."""
|
|
13
|
+
"""Enum defining the types of costs that can be registered."""
|
|
14
14
|
|
|
15
|
-
OTHER =
|
|
16
|
-
TOKEN_INPUT =
|
|
17
|
-
TOKEN_OUTPUT =
|
|
18
|
-
API_CALL =
|
|
19
|
-
STORAGE =
|
|
20
|
-
TIME =
|
|
15
|
+
OTHER = "OTHER"
|
|
16
|
+
TOKEN_INPUT = "TOKEN_INPUT"
|
|
17
|
+
TOKEN_OUTPUT = "TOKEN_OUTPUT"
|
|
18
|
+
API_CALL = "API_CALL"
|
|
19
|
+
STORAGE = "STORAGE"
|
|
20
|
+
TIME = "TIME"
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class CostConfig(BaseModel):
|
|
24
|
+
"""Pydantic model that defines a cost configuration.
|
|
25
|
+
|
|
26
|
+
:param cost_name: Name of the cost (unique identifier in the service).
|
|
27
|
+
:param cost_type: The type/category of the cost.
|
|
28
|
+
:param description: A short description of the cost.
|
|
29
|
+
:param unit: The unit of measurement (e.g. token, call, MB).
|
|
30
|
+
:param rate: The cost per unit (e.g. dollars per token).
|
|
31
|
+
"""
|
|
32
|
+
|
|
33
|
+
cost_name: str
|
|
34
|
+
cost_type: Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]
|
|
35
|
+
description: str | None = None
|
|
36
|
+
unit: str
|
|
37
|
+
rate: float
|
|
21
38
|
|
|
22
39
|
|
|
23
40
|
class CostData(BaseModel):
|
|
24
|
-
"""."""
|
|
41
|
+
"""Data model for cost operations."""
|
|
25
42
|
|
|
26
43
|
cost: float
|
|
27
44
|
mission_id: str
|
|
28
45
|
name: str
|
|
29
|
-
|
|
46
|
+
cost_type: CostType
|
|
30
47
|
unit: str
|
|
48
|
+
rate: float
|
|
49
|
+
setup_version_id: str
|
|
50
|
+
quantity: float
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
class CostServiceError(Exception):
|
|
54
|
+
"""Custom exception for CostService errors."""
|
|
31
55
|
|
|
32
56
|
|
|
33
57
|
class CostStrategy(BaseStrategy, ABC):
|
|
34
58
|
"""Abstract base class for cost strategies."""
|
|
35
59
|
|
|
60
|
+
def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, CostConfig]) -> None:
|
|
61
|
+
"""Initialize the strategy.
|
|
62
|
+
|
|
63
|
+
Args:
|
|
64
|
+
mission_id: The ID of the mission this strategy is associated with
|
|
65
|
+
setup_version_id: The ID of the setup version this strategy is associated with
|
|
66
|
+
config: Configuration dictionary for the strategy
|
|
67
|
+
"""
|
|
68
|
+
super().__init__(mission_id, setup_version_id)
|
|
69
|
+
self.config = config
|
|
70
|
+
|
|
36
71
|
@abstractmethod
|
|
37
|
-
def
|
|
72
|
+
def add(
|
|
73
|
+
self,
|
|
74
|
+
name: str,
|
|
75
|
+
cost_config_name: str,
|
|
76
|
+
quantity: float,
|
|
77
|
+
) -> None:
|
|
38
78
|
"""Register a new cost."""
|
|
39
79
|
|
|
40
|
-
def __post_init__(self, *args, **kwargs) -> None: # noqa: ANN002, ANN003
|
|
41
|
-
"""Allow post init configuration."""
|
|
42
|
-
|
|
43
80
|
@abstractmethod
|
|
44
|
-
def get(
|
|
81
|
+
def get(
|
|
82
|
+
self,
|
|
83
|
+
name: str,
|
|
84
|
+
) -> list[CostData]:
|
|
45
85
|
"""Get a cost."""
|
|
86
|
+
|
|
87
|
+
@abstractmethod
|
|
88
|
+
def get_filtered(
|
|
89
|
+
self,
|
|
90
|
+
names: list[str] | None = None,
|
|
91
|
+
cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
|
|
92
|
+
) -> list[CostData]:
|
|
93
|
+
"""Get filtered costs."""
|
|
@@ -1,30 +1,113 @@
|
|
|
1
1
|
"""Default cost."""
|
|
2
2
|
|
|
3
|
-
import
|
|
4
|
-
from typing import Any
|
|
3
|
+
from typing import Literal
|
|
5
4
|
|
|
6
|
-
from digitalkin.
|
|
7
|
-
|
|
8
|
-
|
|
5
|
+
from digitalkin.logger import logger
|
|
6
|
+
from digitalkin.services.cost.cost_strategy import (
|
|
7
|
+
CostConfig,
|
|
8
|
+
CostData,
|
|
9
|
+
CostServiceError,
|
|
10
|
+
CostStrategy,
|
|
11
|
+
CostType,
|
|
12
|
+
)
|
|
9
13
|
|
|
10
14
|
|
|
11
15
|
class DefaultCost(CostStrategy):
|
|
12
16
|
"""Default cost strategy."""
|
|
13
17
|
|
|
14
|
-
def
|
|
18
|
+
def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, CostConfig]) -> None:
|
|
19
|
+
"""Initialize the strategy.
|
|
20
|
+
|
|
21
|
+
Args:
|
|
22
|
+
mission_id: The ID of the mission this strategy is associated with
|
|
23
|
+
setup_version_id: The ID of the setup version this strategy is associated with
|
|
24
|
+
config: The configuration dictionary for the cost
|
|
25
|
+
"""
|
|
26
|
+
super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
|
|
27
|
+
self.db: dict[str, list[CostData]] = {}
|
|
28
|
+
|
|
29
|
+
def add(
|
|
30
|
+
self,
|
|
31
|
+
name: str,
|
|
32
|
+
cost_config_name: str,
|
|
33
|
+
quantity: float,
|
|
34
|
+
) -> None:
|
|
15
35
|
"""Create a new record in the cost database.
|
|
16
36
|
|
|
37
|
+
Args:
|
|
38
|
+
name: The name of the cost
|
|
39
|
+
cost_config_name: The name of the cost config
|
|
40
|
+
quantity: The quantity of the cost
|
|
41
|
+
|
|
42
|
+
Raises:
|
|
43
|
+
CostServiceError: If the cost data is invalid or if the cost already exists
|
|
44
|
+
"""
|
|
45
|
+
cost_config = self.config.get(cost_config_name)
|
|
46
|
+
if cost_config is None:
|
|
47
|
+
msg = f"Cost config {cost_config_name} not found in the configuration."
|
|
48
|
+
logger.error(msg)
|
|
49
|
+
raise CostServiceError(msg)
|
|
50
|
+
cost_data = CostData.model_validate({
|
|
51
|
+
"name": name,
|
|
52
|
+
"cost": cost_config.rate * quantity,
|
|
53
|
+
"unit": cost_config.unit,
|
|
54
|
+
"cost_type": getattr(CostType, cost_config.cost_type),
|
|
55
|
+
"mission_id": self.mission_id,
|
|
56
|
+
"rate": cost_config.rate,
|
|
57
|
+
"quantity": quantity,
|
|
58
|
+
"setup_version_id": self.setup_version_id,
|
|
59
|
+
})
|
|
60
|
+
if cost_data.mission_id not in self.db:
|
|
61
|
+
self.db[cost_data.mission_id] = []
|
|
62
|
+
if cost_data.name in [cost.name for cost in self.db[cost_data.mission_id]]:
|
|
63
|
+
msg = f"Cost with name {cost_data.name} already exists in mission {cost_data.mission_id}"
|
|
64
|
+
logger.error(msg)
|
|
65
|
+
raise CostServiceError(msg)
|
|
66
|
+
self.db[cost_data.mission_id].append(cost_data)
|
|
67
|
+
|
|
68
|
+
def get(self, name: str) -> list[CostData]:
|
|
69
|
+
"""Get a record from the database.
|
|
70
|
+
|
|
71
|
+
Args:
|
|
72
|
+
name: The name of the cost
|
|
73
|
+
|
|
17
74
|
Returns:
|
|
18
|
-
|
|
75
|
+
list[CostData]: The cost data
|
|
76
|
+
|
|
77
|
+
Raises:
|
|
78
|
+
CostServiceError: If the cost data is invalid or if the cost does not exist
|
|
19
79
|
"""
|
|
20
|
-
|
|
21
|
-
|
|
80
|
+
if self.mission_id not in self.db:
|
|
81
|
+
msg = f"Mission {self.mission_id} not found in the database."
|
|
82
|
+
logger.warning(msg)
|
|
83
|
+
raise CostServiceError(msg)
|
|
84
|
+
|
|
85
|
+
return [cost for cost in self.db[self.mission_id] if cost.name == name] or []
|
|
22
86
|
|
|
23
|
-
def
|
|
87
|
+
def get_filtered(
|
|
88
|
+
self,
|
|
89
|
+
names: list[str] | None = None,
|
|
90
|
+
cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
|
|
91
|
+
) -> list[CostData]:
|
|
24
92
|
"""Get records from the database.
|
|
25
93
|
|
|
94
|
+
Args:
|
|
95
|
+
names: The names of the costs
|
|
96
|
+
cost_types: The types of the costs
|
|
97
|
+
|
|
26
98
|
Returns:
|
|
27
99
|
list[CostData]: The list of records
|
|
100
|
+
|
|
101
|
+
Raises:
|
|
102
|
+
CostServiceError: If the cost data is invalid or if the cost does not exist
|
|
28
103
|
"""
|
|
29
|
-
|
|
30
|
-
|
|
104
|
+
if self.mission_id not in self.db:
|
|
105
|
+
msg = f"Mission {self.mission_id} not found in the database."
|
|
106
|
+
logger.warning(msg)
|
|
107
|
+
raise CostServiceError(msg)
|
|
108
|
+
|
|
109
|
+
return [
|
|
110
|
+
cost
|
|
111
|
+
for cost in self.db[self.mission_id]
|
|
112
|
+
if (names and cost.name in names) or (cost_types and cost.cost_type in cost_types)
|
|
113
|
+
]
|
|
@@ -1,81 +1,170 @@
|
|
|
1
|
-
"""This module implements the
|
|
1
|
+
"""This module implements the gRPC Cost strategy."""
|
|
2
2
|
|
|
3
|
-
import
|
|
4
|
-
from
|
|
3
|
+
from collections.abc import Generator
|
|
4
|
+
from contextlib import contextmanager
|
|
5
|
+
from typing import Any, Literal
|
|
5
6
|
|
|
6
7
|
from digitalkin_proto.digitalkin.cost.v1 import cost_pb2, cost_service_pb2_grpc
|
|
7
8
|
from google.protobuf import json_format
|
|
8
|
-
from pydantic import ValidationError
|
|
9
9
|
|
|
10
|
+
from digitalkin.grpc_servers.utils.exceptions import ServerError
|
|
10
11
|
from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
|
|
11
|
-
from digitalkin.grpc_servers.utils.models import
|
|
12
|
-
from digitalkin.
|
|
13
|
-
|
|
14
|
-
|
|
12
|
+
from digitalkin.grpc_servers.utils.models import ClientConfig
|
|
13
|
+
from digitalkin.logger import logger
|
|
14
|
+
from digitalkin.services.cost.cost_strategy import (
|
|
15
|
+
CostConfig,
|
|
16
|
+
CostData,
|
|
17
|
+
CostServiceError,
|
|
18
|
+
CostStrategy,
|
|
19
|
+
CostType,
|
|
20
|
+
)
|
|
15
21
|
|
|
16
22
|
|
|
17
23
|
class GrpcCost(CostStrategy, GrpcClientWrapper):
|
|
18
24
|
"""This class implements the default Cost strategy."""
|
|
19
25
|
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
)
|
|
25
|
-
response = self.exec_grpc_query("GetCostsByName", request)
|
|
26
|
-
return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
|
|
27
|
-
|
|
28
|
-
def _get_costs_by_mission(self, cost_dict: dict[str, Any]) -> list[CostData]:
|
|
29
|
-
request = cost_pb2.GetCostsByMissionRequest(mission_id=cost_dict["mission_id"])
|
|
30
|
-
response = self.exec_grpc_query("GetCostsByMission", request)
|
|
31
|
-
return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
|
|
32
|
-
|
|
33
|
-
def _get_costs_by_type(self, cost_dict: dict[str, Any]) -> list[CostData]:
|
|
34
|
-
request = cost_pb2.GetCostsByTypeRequest(
|
|
35
|
-
mission_id=cost_dict["mission_id"],
|
|
36
|
-
type=cost_dict["type"],
|
|
37
|
-
)
|
|
38
|
-
response = self.exec_grpc_query("GetCostsBytype", request)
|
|
39
|
-
return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
|
|
40
|
-
|
|
41
|
-
def __init__(self, mission_id: str, config: ServerConfig) -> None:
|
|
42
|
-
"""Initialize the cost."""
|
|
43
|
-
super().__init__(mission_id)
|
|
44
|
-
channel = self._init_channel(config)
|
|
45
|
-
self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
|
|
46
|
-
logger.info("Channel client 'Cost' initialized succesfully")
|
|
26
|
+
@staticmethod
|
|
27
|
+
@contextmanager
|
|
28
|
+
def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
|
|
29
|
+
"""Context manager for consistent gRPC error handling.
|
|
47
30
|
|
|
48
|
-
|
|
49
|
-
|
|
31
|
+
Yields:
|
|
32
|
+
Allow error handling in context.
|
|
50
33
|
|
|
51
|
-
|
|
52
|
-
|
|
34
|
+
Args:
|
|
35
|
+
operation: Description of the operation being performed.
|
|
53
36
|
|
|
54
|
-
|
|
55
|
-
|
|
37
|
+
Raises:
|
|
38
|
+
ValueError: Error with the model validation.
|
|
39
|
+
ServerError: from gRPC Client.
|
|
40
|
+
CostServiceError: Unexpected error.
|
|
56
41
|
"""
|
|
57
42
|
try:
|
|
58
|
-
|
|
59
|
-
except
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
43
|
+
yield
|
|
44
|
+
except CostServiceError as e:
|
|
45
|
+
msg = f"CostServiceError in {operation}: {e}"
|
|
46
|
+
logger.exception(msg)
|
|
47
|
+
raise CostServiceError(msg) from e
|
|
48
|
+
except ServerError as e:
|
|
49
|
+
msg = f"gRPC {operation} failed: {e}"
|
|
50
|
+
logger.exception(msg)
|
|
51
|
+
raise ServerError(msg) from e
|
|
52
|
+
except Exception as e:
|
|
53
|
+
msg = f"Unexpected error in {operation}"
|
|
54
|
+
logger.exception(msg)
|
|
55
|
+
raise CostServiceError(msg) from e
|
|
56
|
+
|
|
57
|
+
def __init__(
|
|
58
|
+
self,
|
|
59
|
+
mission_id: str,
|
|
60
|
+
setup_version_id: str,
|
|
61
|
+
config: dict[str, CostConfig],
|
|
62
|
+
client_config: ClientConfig,
|
|
63
|
+
) -> None:
|
|
64
|
+
"""Initialize the cost."""
|
|
65
|
+
super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
|
|
66
|
+
channel = self._init_channel(client_config)
|
|
67
|
+
self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
|
|
68
|
+
logger.debug("Channel client 'Cost' initialized succesfully")
|
|
69
|
+
|
|
70
|
+
def add(
|
|
71
|
+
self,
|
|
72
|
+
name: str,
|
|
73
|
+
cost_config_name: str,
|
|
74
|
+
quantity: float,
|
|
75
|
+
) -> None:
|
|
76
|
+
"""Create a new record in the cost database.
|
|
77
|
+
|
|
78
|
+
Args:
|
|
79
|
+
name: The name of the cost
|
|
80
|
+
cost_config_name: The name of the cost config
|
|
81
|
+
quantity: The quantity of the cost
|
|
65
82
|
|
|
66
|
-
|
|
67
|
-
|
|
83
|
+
Raises:
|
|
84
|
+
CostServiceError: If the cost config is invalid
|
|
85
|
+
"""
|
|
86
|
+
with self._handle_grpc_errors("AddCost"):
|
|
87
|
+
cost_config = self.config.get(cost_config_name)
|
|
88
|
+
if cost_config is None:
|
|
89
|
+
msg = f"Cost config {cost_config_name} not found in the configuration."
|
|
90
|
+
logger.error(msg)
|
|
91
|
+
raise CostServiceError(msg)
|
|
92
|
+
valid_data = CostData.model_validate({
|
|
93
|
+
"name": name,
|
|
94
|
+
"cost": cost_config.rate * quantity,
|
|
95
|
+
"unit": cost_config.unit,
|
|
96
|
+
"cost_type": CostType[cost_config.cost_type],
|
|
97
|
+
"mission_id": self.mission_id,
|
|
98
|
+
"rate": cost_config.rate,
|
|
99
|
+
"quantity": quantity,
|
|
100
|
+
"setup_version_id": self.setup_version_id,
|
|
101
|
+
})
|
|
102
|
+
request = cost_pb2.AddCostRequest(
|
|
103
|
+
cost=valid_data.cost,
|
|
104
|
+
name=valid_data.name,
|
|
105
|
+
unit=valid_data.unit,
|
|
106
|
+
cost_type=valid_data.cost_type.name,
|
|
107
|
+
mission_id=valid_data.mission_id,
|
|
108
|
+
rate=valid_data.rate,
|
|
109
|
+
quantity=valid_data.quantity,
|
|
110
|
+
setup_version_id=valid_data.setup_version_id,
|
|
111
|
+
)
|
|
112
|
+
self.exec_grpc_query("AddCost", request)
|
|
113
|
+
logger.debug("Cost added with cost_dict: %s", valid_data.model_dump())
|
|
114
|
+
|
|
115
|
+
def get(self, name: str) -> list[CostData]:
|
|
116
|
+
"""Get a record from the database.
|
|
117
|
+
|
|
118
|
+
Args:
|
|
119
|
+
name: The name of the cost
|
|
68
120
|
|
|
69
|
-
|
|
70
|
-
|
|
121
|
+
Returns:
|
|
122
|
+
CostData: The cost data
|
|
123
|
+
"""
|
|
124
|
+
with self._handle_grpc_errors("GetCost"):
|
|
125
|
+
request = cost_pb2.GetCostRequest(name=name, mission_id=self.mission_id)
|
|
126
|
+
response: cost_pb2.GetCostResponse = self.exec_grpc_query("GetCost", request)
|
|
127
|
+
cost_data_list = [
|
|
128
|
+
json_format.MessageToDict(
|
|
129
|
+
cost,
|
|
130
|
+
preserving_proto_field_name=True,
|
|
131
|
+
always_print_fields_with_no_presence=True,
|
|
132
|
+
)
|
|
133
|
+
for cost in response.costs
|
|
134
|
+
]
|
|
135
|
+
logger.debug("Costs retrieved with cost_dict: %s", cost_data_list)
|
|
136
|
+
return [CostData.model_validate(cost_data) for cost_data in cost_data_list]
|
|
137
|
+
|
|
138
|
+
def get_filtered(
|
|
139
|
+
self,
|
|
140
|
+
names: list[str] | None = None,
|
|
141
|
+
cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
|
|
142
|
+
) -> list[CostData]:
|
|
143
|
+
"""Get a list of records from the database.
|
|
144
|
+
|
|
145
|
+
Args:
|
|
146
|
+
names: The names of the costs
|
|
147
|
+
cost_types: The types of the costs
|
|
71
148
|
|
|
72
149
|
Returns:
|
|
73
|
-
list[CostData]: The
|
|
150
|
+
list[CostData]: The cost data
|
|
74
151
|
"""
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
152
|
+
with self._handle_grpc_errors("GetCosts"):
|
|
153
|
+
request = cost_pb2.GetCostsRequest(
|
|
154
|
+
mission_id=self.mission_id,
|
|
155
|
+
filter=cost_pb2.CostFilter(
|
|
156
|
+
names=names or [],
|
|
157
|
+
cost_types=cost_types or [],
|
|
158
|
+
),
|
|
159
|
+
)
|
|
160
|
+
response: cost_pb2.GetCostsResponse = self.exec_grpc_query("GetCosts", request)
|
|
161
|
+
cost_data_list = [
|
|
162
|
+
json_format.MessageToDict(
|
|
163
|
+
cost,
|
|
164
|
+
preserving_proto_field_name=True,
|
|
165
|
+
always_print_fields_with_no_presence=True,
|
|
166
|
+
)
|
|
167
|
+
for cost in response.costs
|
|
168
|
+
]
|
|
169
|
+
logger.debug("Filtered costs retrieved with cost_dict: %s", cost_data_list)
|
|
170
|
+
return [CostData.model_validate(cost_data) for cost_data in cost_data_list]
|
|
@@ -1,10 +1,10 @@
|
|
|
1
1
|
"""Default filesystem."""
|
|
2
2
|
|
|
3
|
-
import logging
|
|
4
3
|
import os
|
|
5
4
|
import tempfile
|
|
6
5
|
from pathlib import Path
|
|
7
6
|
|
|
7
|
+
from digitalkin.logger import logger
|
|
8
8
|
from digitalkin.services.filesystem.filesystem_strategy import (
|
|
9
9
|
FilesystemData,
|
|
10
10
|
FilesystemServiceError,
|
|
@@ -12,20 +12,19 @@ from digitalkin.services.filesystem.filesystem_strategy import (
|
|
|
12
12
|
FileType,
|
|
13
13
|
)
|
|
14
14
|
|
|
15
|
-
logger = logging.getLogger(__name__)
|
|
16
|
-
|
|
17
15
|
|
|
18
16
|
class DefaultFilesystem(FilesystemStrategy):
|
|
19
17
|
"""Default state filesystem strategy."""
|
|
20
18
|
|
|
21
|
-
def __init__(self, mission_id: str, config: dict[str, str]) -> None:
|
|
19
|
+
def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, str]) -> None:
|
|
22
20
|
"""Initialize the default filesystem strategy.
|
|
23
21
|
|
|
24
22
|
Args:
|
|
25
23
|
mission_id: The ID of the mission this strategy is associated with
|
|
24
|
+
setup_version_id: The ID of the setup version this strategy is associated with
|
|
26
25
|
config: A dictionary mapping names to Pydantic model classes
|
|
27
26
|
"""
|
|
28
|
-
super().__init__(mission_id, config)
|
|
27
|
+
super().__init__(mission_id, setup_version_id, config)
|
|
29
28
|
self.temp_root: str = self.config.get("temp_root", "") or tempfile.gettempdir()
|
|
30
29
|
os.makedirs(self.temp_root, exist_ok=True)
|
|
31
30
|
self.db: dict[str, FilesystemData] = {}
|
|
@@ -176,7 +175,7 @@ class DefaultFilesystem(FilesystemStrategy):
|
|
|
176
175
|
try:
|
|
177
176
|
os.remove(file_path)
|
|
178
177
|
del self.db[name]
|
|
179
|
-
logger.
|
|
178
|
+
logger.debug("File %s successfully deleted.", name)
|
|
180
179
|
|
|
181
180
|
except OSError:
|
|
182
181
|
msg = f"Error deleting file {name} from filesystem"
|
|
@@ -35,14 +35,15 @@ class FilesystemData(BaseModel):
|
|
|
35
35
|
class FilesystemStrategy(BaseStrategy, ABC):
|
|
36
36
|
"""Abstract base class for filesystem strategies."""
|
|
37
37
|
|
|
38
|
-
def __init__(self, mission_id: str, config: dict[str, str]) -> None:
|
|
38
|
+
def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, str]) -> None:
|
|
39
39
|
"""Initialize the strategy.
|
|
40
40
|
|
|
41
41
|
Args:
|
|
42
42
|
mission_id: The ID of the mission this strategy is associated with
|
|
43
|
+
setup_version_id: The ID of the setup version this strategy is associated with
|
|
43
44
|
config: configuration dictionary for the filesystem strategy
|
|
44
45
|
"""
|
|
45
|
-
super().__init__(mission_id)
|
|
46
|
+
super().__init__(mission_id, setup_version_id)
|
|
46
47
|
self.config: dict[str, str] = config
|
|
47
48
|
|
|
48
49
|
@abstractmethod
|
|
@@ -1,16 +1,21 @@
|
|
|
1
1
|
"""Grpc filesystem."""
|
|
2
2
|
|
|
3
|
-
import logging
|
|
4
3
|
from collections.abc import Generator
|
|
5
4
|
from contextlib import contextmanager
|
|
6
5
|
from typing import Any
|
|
7
6
|
|
|
8
|
-
from digitalkin_proto.digitalkin.filesystem.v2 import
|
|
9
|
-
|
|
7
|
+
from digitalkin_proto.digitalkin.filesystem.v2 import (
|
|
8
|
+
filesystem_pb2,
|
|
9
|
+
filesystem_service_pb2_grpc,
|
|
10
|
+
)
|
|
11
|
+
from digitalkin_proto.digitalkin.filesystem.v2.filesystem_pb2 import (
|
|
12
|
+
FileType as FileTypeProto,
|
|
13
|
+
)
|
|
10
14
|
|
|
11
15
|
from digitalkin.grpc_servers.utils.exceptions import ServerError
|
|
12
16
|
from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
|
|
13
17
|
from digitalkin.grpc_servers.utils.models import ClientConfig
|
|
18
|
+
from digitalkin.logger import logger
|
|
14
19
|
from digitalkin.services.filesystem.filesystem_strategy import (
|
|
15
20
|
FilesystemData,
|
|
16
21
|
FilesystemServiceError,
|
|
@@ -18,32 +23,10 @@ from digitalkin.services.filesystem.filesystem_strategy import (
|
|
|
18
23
|
FileType,
|
|
19
24
|
)
|
|
20
25
|
|
|
21
|
-
logger = logging.getLogger(__name__)
|
|
22
|
-
|
|
23
26
|
|
|
24
27
|
class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
|
|
25
28
|
"""Default state filesystem strategy."""
|
|
26
29
|
|
|
27
|
-
def __init__(
|
|
28
|
-
self,
|
|
29
|
-
mission_id: str,
|
|
30
|
-
config: dict[str, str],
|
|
31
|
-
client_config: ClientConfig,
|
|
32
|
-
**kwargs, # noqa: ANN003, ARG002
|
|
33
|
-
) -> None:
|
|
34
|
-
"""Initialize the default filesystem strategy.
|
|
35
|
-
|
|
36
|
-
Args:
|
|
37
|
-
mission_id: The ID of the mission this strategy is associated with
|
|
38
|
-
config: A dictionary mapping names to Pydantic model classes
|
|
39
|
-
client_config: The client configuration object
|
|
40
|
-
kwargs: other optional arguments to pass to the parent class constructor
|
|
41
|
-
"""
|
|
42
|
-
super().__init__(mission_id, config)
|
|
43
|
-
channel = self._init_channel(client_config)
|
|
44
|
-
self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
|
|
45
|
-
logger.info("Channel client 'Filesystem' initialized succesfully")
|
|
46
|
-
|
|
47
30
|
@staticmethod
|
|
48
31
|
@contextmanager
|
|
49
32
|
def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
|
|
@@ -56,7 +39,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
|
|
|
56
39
|
operation: Description of the operation being performed.
|
|
57
40
|
|
|
58
41
|
Raises:
|
|
59
|
-
ValueError: Error
|
|
42
|
+
ValueError: Error with the model validation.
|
|
60
43
|
ServerError: from gRPC Client.
|
|
61
44
|
FilesystemServiceError: Filesystem service internal.
|
|
62
45
|
"""
|
|
@@ -71,6 +54,28 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
|
|
|
71
54
|
logger.exception(msg)
|
|
72
55
|
raise FilesystemServiceError(msg) from e
|
|
73
56
|
|
|
57
|
+
def __init__(
|
|
58
|
+
self,
|
|
59
|
+
mission_id: str,
|
|
60
|
+
setup_version_id: str,
|
|
61
|
+
config: dict[str, str],
|
|
62
|
+
client_config: ClientConfig,
|
|
63
|
+
**kwargs, # noqa: ANN003, ARG002
|
|
64
|
+
) -> None:
|
|
65
|
+
"""Initialize the default filesystem strategy.
|
|
66
|
+
|
|
67
|
+
Args:
|
|
68
|
+
mission_id: The ID of the mission this strategy is associated with
|
|
69
|
+
setup_version_id: The ID of the setup version this strategy is associated with
|
|
70
|
+
config: A dictionary mapping names to Pydantic model classes
|
|
71
|
+
client_config: The server configuration object
|
|
72
|
+
kwargs: other optional arguments to pass to the parent class constructor
|
|
73
|
+
"""
|
|
74
|
+
super().__init__(mission_id, setup_version_id, config)
|
|
75
|
+
channel = self._init_channel(client_config)
|
|
76
|
+
self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
|
|
77
|
+
logger.info("Channel client 'Filesystem' initialized succesfully")
|
|
78
|
+
|
|
74
79
|
def upload(self, content: bytes, name: str, file_type: FileType) -> FilesystemData:
|
|
75
80
|
"""Create a new file in the file system.
|
|
76
81
|
|
|
@@ -5,7 +5,7 @@ from typing import Any, ClassVar
|
|
|
5
5
|
from pydantic import BaseModel, Field, PrivateAttr
|
|
6
6
|
|
|
7
7
|
from digitalkin.services.agent import AgentStrategy, DefaultAgent
|
|
8
|
-
from digitalkin.services.cost import CostStrategy, DefaultCost
|
|
8
|
+
from digitalkin.services.cost import CostStrategy, DefaultCost, GrpcCost
|
|
9
9
|
from digitalkin.services.filesystem import DefaultFilesystem, FilesystemStrategy, GrpcFilesystem
|
|
10
10
|
from digitalkin.services.identity import DefaultIdentity, IdentityStrategy
|
|
11
11
|
from digitalkin.services.registry import DefaultRegistry, RegistryStrategy
|
|
@@ -30,7 +30,7 @@ class ServicesConfig(BaseModel):
|
|
|
30
30
|
)
|
|
31
31
|
_config_storage: dict[str, Any | None] = PrivateAttr(default_factory=dict)
|
|
32
32
|
_cost: ServicesStrategy[CostStrategy] = PrivateAttr(
|
|
33
|
-
default_factory=lambda: ServicesStrategy(local=DefaultCost, remote=
|
|
33
|
+
default_factory=lambda: ServicesStrategy(local=DefaultCost, remote=GrpcCost)
|
|
34
34
|
)
|
|
35
35
|
_config_cost: dict[str, Any | None] = PrivateAttr(default_factory=dict)
|
|
36
36
|
_snapshot: ServicesStrategy[SnapshotStrategy] = PrivateAttr(
|
|
@@ -111,12 +111,13 @@ class ServicesConfig(BaseModel):
|
|
|
111
111
|
"""
|
|
112
112
|
return getattr(self, f"_config_{name}", {})
|
|
113
113
|
|
|
114
|
-
def init_strategy(self, name: str, mission_id: str) -> ServicesStrategy:
|
|
114
|
+
def init_strategy(self, name: str, mission_id: str, setup_version_id: str) -> ServicesStrategy:
|
|
115
115
|
"""Initialize a specific strategy.
|
|
116
116
|
|
|
117
117
|
Args:
|
|
118
118
|
name: The name of the strategy to initialize
|
|
119
119
|
mission_id: The ID of the mission this strategy is associated with
|
|
120
|
+
setup_version_id: The setup version ID for the strategy
|
|
120
121
|
|
|
121
122
|
Returns:
|
|
122
123
|
The initialized strategy instance
|
|
@@ -129,8 +130,8 @@ class ServicesConfig(BaseModel):
|
|
|
129
130
|
msg = f"Strategy {name} not found in ServicesConfig."
|
|
130
131
|
raise ValueError(msg)
|
|
131
132
|
|
|
132
|
-
# Instantiate the strategy with the mission ID and configuration
|
|
133
|
-
return strategy_type(mission_id, **self.get_strategy_config(name) or {})
|
|
133
|
+
# Instantiate the strategy with the mission ID, setup version ID, and configuration
|
|
134
|
+
return strategy_type(mission_id, setup_version_id, **self.get_strategy_config(name) or {})
|
|
134
135
|
|
|
135
136
|
@property
|
|
136
137
|
def storage(self) -> type[StorageStrategy]:
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""This module is responsible for handling the setup service."""
|