digitalkin 0.2.12__py3-none-any.whl → 0.2.14__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.
Files changed (42) hide show
  1. digitalkin/__version__.py +1 -1
  2. digitalkin/grpc_servers/_base_server.py +15 -17
  3. digitalkin/grpc_servers/module_server.py +9 -10
  4. digitalkin/grpc_servers/module_servicer.py +199 -85
  5. digitalkin/grpc_servers/registry_server.py +3 -6
  6. digitalkin/grpc_servers/registry_servicer.py +18 -19
  7. digitalkin/grpc_servers/utils/exceptions.py +4 -0
  8. digitalkin/grpc_servers/utils/grpc_client_wrapper.py +3 -5
  9. digitalkin/logger.py +45 -1
  10. digitalkin/models/module/__init__.py +2 -1
  11. digitalkin/models/module/module.py +1 -0
  12. digitalkin/models/module/module_types.py +1 -0
  13. digitalkin/modules/_base_module.py +124 -7
  14. digitalkin/modules/archetype_module.py +11 -1
  15. digitalkin/modules/job_manager/base_job_manager.py +181 -0
  16. digitalkin/modules/job_manager/job_manager_models.py +44 -0
  17. digitalkin/modules/job_manager/single_job_manager.py +285 -0
  18. digitalkin/modules/job_manager/taskiq_broker.py +214 -0
  19. digitalkin/modules/job_manager/taskiq_job_manager.py +286 -0
  20. digitalkin/modules/tool_module.py +2 -1
  21. digitalkin/modules/trigger_module.py +3 -1
  22. digitalkin/services/cost/default_cost.py +8 -4
  23. digitalkin/services/cost/grpc_cost.py +15 -7
  24. digitalkin/services/filesystem/default_filesystem.py +2 -4
  25. digitalkin/services/filesystem/grpc_filesystem.py +8 -5
  26. digitalkin/services/setup/__init__.py +1 -0
  27. digitalkin/services/setup/default_setup.py +10 -12
  28. digitalkin/services/setup/grpc_setup.py +8 -10
  29. digitalkin/services/storage/default_storage.py +11 -5
  30. digitalkin/services/storage/grpc_storage.py +23 -8
  31. digitalkin/utils/arg_parser.py +5 -48
  32. digitalkin/utils/development_mode_action.py +51 -0
  33. {digitalkin-0.2.12.dist-info → digitalkin-0.2.14.dist-info}/METADATA +46 -15
  34. {digitalkin-0.2.12.dist-info → digitalkin-0.2.14.dist-info}/RECORD +41 -34
  35. {digitalkin-0.2.12.dist-info → digitalkin-0.2.14.dist-info}/WHEEL +1 -1
  36. modules/cpu_intensive_module.py +281 -0
  37. modules/minimal_llm_module.py +240 -58
  38. modules/storage_module.py +5 -6
  39. modules/text_transform_module.py +1 -1
  40. digitalkin/modules/job_manager.py +0 -177
  41. {digitalkin-0.2.12.dist-info → digitalkin-0.2.14.dist-info}/licenses/LICENSE +0 -0
  42. {digitalkin-0.2.12.dist-info → digitalkin-0.2.14.dist-info}/top_level.txt +0 -0
@@ -6,7 +6,7 @@ from typing import Any, ClassVar
6
6
 
7
7
  from pydantic import BaseModel
8
8
 
9
- from digitalkin.grpc_servers.utils.models import SecurityMode, ClientConfig, ServerMode
9
+ from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerMode
10
10
  from digitalkin.modules._base_module import BaseModule
11
11
  from digitalkin.services.setup.setup_strategy import SetupData
12
12
  from digitalkin.services.storage.storage_strategy import DataType, StorageRecord
@@ -1,177 +0,0 @@
1
- """Background module manager."""
2
-
3
- import asyncio
4
- import uuid
5
- from argparse import ArgumentParser, Namespace
6
- from collections.abc import Callable, Coroutine
7
- from typing import Any
8
-
9
- from digitalkin.logger import logger
10
- from digitalkin.models import ModuleStatus
11
- from digitalkin.models.module import InputModelT, OutputModelT, SecretModelT, SetupModelT
12
- from digitalkin.modules._base_module import BaseModule
13
- from digitalkin.services.services_config import ServicesConfig
14
- from digitalkin.services.services_models import ServicesMode
15
- from digitalkin.utils.arg_parser import ArgParser, DevelopmentModeMappingAction
16
-
17
-
18
- class JobManager(ArgParser):
19
- """Background module manager."""
20
-
21
- args: Namespace
22
-
23
- @staticmethod
24
- async def _job_specific_callback(
25
- callback: Callable[[str, OutputModelT], Coroutine[Any, Any, None]], job_id: str
26
- ) -> Callable[[OutputModelT], Coroutine[Any, Any, None]]:
27
- """Return a callback function for the job.
28
-
29
- Args:
30
- callback: Callback function to be called when the job is done
31
- job_id: Identifiant du module
32
-
33
- Returns:
34
- Callable: Callback function
35
- """
36
-
37
- def callback_wrapper(output_data: OutputModelT) -> Coroutine[Any, Any, None]:
38
- """Wrapper for the callback function.
39
-
40
- Args:
41
- output_data: Output data of the job
42
-
43
- Returns:
44
- Coroutine: Callback function
45
- """
46
- return callback(job_id, output_data)
47
-
48
- return callback_wrapper
49
-
50
- def _add_parser_args(self, parser: ArgumentParser) -> None:
51
- super()._add_parser_args(parser)
52
- parser.add_argument(
53
- "-d",
54
- "--dev-mode",
55
- env_var="SERVICE_MODE",
56
- choices=ServicesMode.__members__,
57
- default="local",
58
- action=DevelopmentModeMappingAction,
59
- dest="services_mode",
60
- help="Define Module Service configurations for endpoints",
61
- )
62
-
63
- def __init__(self, module_class: type[BaseModule]) -> None:
64
- """Initialize the job manager."""
65
- self.module_class = module_class
66
- self.modules: dict[str, BaseModule] = {}
67
- self._lock = asyncio.Lock()
68
- super().__init__()
69
-
70
- services_config = ServicesConfig(
71
- services_config_strategies=self.module_class.services_config_strategies,
72
- services_config_params=self.module_class.services_config_params,
73
- mode=self.args.services_mode,
74
- )
75
- setattr(self.module_class, "services_config", services_config)
76
-
77
- async def create_job( # noqa: D417
78
- self,
79
- input_data: InputModelT,
80
- setup_data: SetupModelT,
81
- mission_id: str,
82
- setup_version_id: str,
83
- callback: Callable[[str, OutputModelT], Coroutine[Any, Any, None]],
84
- ) -> tuple[str, BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT]]: # type: ignore
85
- """Start new module job in background (asyncio).
86
-
87
- Args:
88
- module_class: Classe du module à instancier
89
- *args: Arguments à passer au constructeur du module
90
- **kwargs: Arguments à passer au constructeur du module
91
-
92
- Returns:
93
- str: job_id of the module entity
94
- """
95
- job_id = str(uuid.uuid4())
96
- """TODO: check uniqueness of the job_id"""
97
- # Création et démarrage du module
98
- module = self.module_class(job_id, mission_id=mission_id, setup_version_id=setup_version_id)
99
- self.modules[job_id] = module
100
- try:
101
- await module.start(input_data, setup_data, await JobManager._job_specific_callback(callback, job_id))
102
- logger.info("Module %s (%s) started successfully", job_id, module.name)
103
- except Exception:
104
- # En cas d'erreur, supprimer le module du gestionnaire
105
- del self.modules[job_id]
106
- logger.exception("Échec du démarrage du module %s: %s", job_id)
107
- raise
108
- else:
109
- return job_id, module
110
-
111
- async def stop_module(self, job_id: str) -> bool:
112
- """Arrête un module en cours d'exécution.
113
-
114
- Args:
115
- job_id: Identifiant du module à arrêter
116
-
117
- Returns:
118
- True si le module a été arrêté, False s'il n'existe pas.
119
- """
120
- async with self._lock:
121
- module = self.modules.get(job_id)
122
- if not module:
123
- logger.warning(f"Module {job_id} introuvable")
124
- return False
125
- try:
126
- await module.stop()
127
- logger.info(f"Module {job_id} ({module.name}) arrêté avec succès")
128
- except Exception as e:
129
- logger.error(f"Erreur lors de l'arrêt du module {job_id}: {e}")
130
- raise
131
- else:
132
- return True
133
-
134
- def get_module_status(self, job_id: str) -> ModuleStatus | None:
135
- """Obtient le statut d'un module.
136
-
137
- Args:
138
- job_id: Identifiant du module
139
-
140
- Returns:
141
- Le statut du module ou None si le module n'existe pas.
142
- """
143
- module = self.modules.get(job_id)
144
- return module.status if module else None
145
-
146
- def get_module(self, job_id: str) -> BaseModule | None:
147
- """Récupère une référence au module.
148
-
149
- Args:
150
- job_id: Identifiant du module
151
-
152
- Returns:
153
- Le module ou None s'il n'existe pas.
154
- """
155
- return self.modules.get(job_id)
156
-
157
- async def stop_all_modules(self) -> None:
158
- """Arrête tous les modules en cours d'exécution."""
159
- async with self._lock:
160
- stop_tasks = [self.stop_module(job_id) for job_id in list(self.modules.keys())]
161
- if stop_tasks:
162
- await asyncio.gather(*stop_tasks, return_exceptions=True)
163
-
164
- def list_modules(self) -> dict[str, dict[str, Any]]:
165
- """Liste tous les modules avec leur statut.
166
-
167
- Returns:
168
- Dictionnaire des modules avec leurs informations.
169
- """
170
- return {
171
- job_id: {
172
- "name": module.name,
173
- "status": module.status,
174
- "class": module.__class__.__name__,
175
- }
176
- for job_id, module in self.modules.items()
177
- }