Coverage for src/backend/InvenTree/machine/registry.py: 54%
286 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 17:47 +0000
1"""Machine registry."""
3import functools
4from typing import Any, Optional, cast
5from uuid import UUID
7from django.db.utils import IntegrityError, OperationalError, ProgrammingError
9import structlog
11import InvenTree.cache
12from common.settings import get_global_setting, set_global_setting
13from InvenTree.exceptions import log_error
14from InvenTree.helpers_mixin import get_shared_class_instance_state_mixin
15from machine.machine_type import BaseDriver, BaseMachineType
17logger = structlog.get_logger('inventree')
20def machine_registry_entrypoint(
21 check_reload: bool = True, check_ready: bool = True, default_value: Any = None
22) -> Any:
23 """Decorator for any method which should be registered as a machine registry entrypoint.
25 This decorator ensures that the plugin registry is up-to-date,
26 and reloads the machine registry if necessary.
27 """
29 def decorator(method):
30 """Internal decorator for the machine registry entrypoint."""
32 @functools.wraps(method)
33 def wrapper(self, *args, **kwargs):
34 """Wrapper function to ensure the machine registry is up-to-date."""
35 # Ensure the plugin registry is up-to-date
36 from plugin import registry as plg_registry
38 logger.debug("machine_registry_entrypoint: '%s'", method.__name__)
40 if check_ready and not self.ready: 40 ↛ 41line 40 didn't jump to line 41 because the condition on line 40 was never true
41 logger.warning(
42 "Machine registry is not ready - cannot call method '%s'",
43 method.__name__,
44 )
46 return default_value
48 do_reload = False
50 if InvenTree.cache.get_session_cache('machine_registry_checked'):
51 # Short circuit if we have already checked within this session
52 pass
54 elif not getattr(self, '__checking_reload', False):
55 # Avoid recursive reloads
56 do_reload = True
57 self.__checking_reload = True
59 if check_reload:
60 if plg_registry.check_reload(): 60 ↛ 62line 60 didn't jump to line 62 because the condition on line 60 was never true
61 # The plugin registry changed - update the machine registry too
62 logger.info(
63 'Plugin registry changed - reloading machine registry'
64 )
65 self.reload_machines()
67 else:
68 # Check if the machine registry needs to be reloaded
69 self._check_reload()
71 self.__checking_reload = False
73 InvenTree.cache.set_session_cache('machine_registry_checked', True)
75 # Call the original method
76 try:
77 result = method(self, *args, **kwargs)
78 except Exception as e:
79 log_error(method.__name__, scope='machine_registry')
80 result = default_value
81 raise e
82 finally:
83 # If we reloaded the registry, we need to update the registry hash
84 if do_reload:
85 self._update_registry_hash()
87 return result
89 return wrapper
91 return decorator
94class MachineRegistry(
95 get_shared_class_instance_state_mixin(lambda _x: 'machine:registry')
96):
97 """Machine registry class."""
99 def __init__(self) -> None:
100 """Initialize machine registry.
102 Set up all needed references for internal and external states.
103 """
104 self.machine_types: dict[str, type[BaseMachineType]] = {}
105 self.drivers: dict[str, type[BaseDriver]] = {}
106 self.driver_instances: dict[str, BaseDriver] = {}
107 self.machines: dict[str, BaseMachineType] = {}
109 self.base_drivers: list[type[BaseDriver]] = []
111 self.ready: bool = False
113 # Keep an internal hash of the machine registry state
114 self._hash = None
116 @property
117 def errors(self) -> list[str | Exception]:
118 """List of registry errors."""
119 return cast(list[str | Exception], self.get_shared_state('errors', []))
121 @property
122 def is_ready(self) -> bool:
123 """Check if the machine registry is ready."""
124 return self.ready
126 def handle_error(self, error: Exception | str):
127 """Helper function for capturing errors with the machine registry."""
128 if error not in self.errors:
129 self.set_shared_state('errors', [*self.errors, error])
131 @machine_registry_entrypoint(check_reload=False, check_ready=False)
132 def initialize(self, main: bool = False):
133 """Initialize the machine registry."""
134 self.ready = True
135 self.reload_machines(main=main)
137 def discover_machine_types(self):
138 """Discovers all machine types by discovering all plugins which implement the Machine mixin class."""
139 from plugin import PluginMixinEnum
140 from plugin.registry import registry as plugin_registry
142 logger.debug('Collecting machine types')
144 machine_types: dict[str, type[BaseMachineType]] = {}
145 base_drivers: list[type[BaseDriver]] = []
147 for plugin in plugin_registry.with_mixin(PluginMixinEnum.MACHINE):
148 try:
149 for machine_type in plugin.get_machine_types():
150 if not issubclass(machine_type, BaseMachineType): 150 ↛ 151line 150 didn't jump to line 151 because the condition on line 150 was never true
151 logger.error(
152 'INVE-E12: Plugin %s returned invalid machine type',
153 plugin.slug,
154 )
155 continue
157 try:
158 machine_type.validate()
159 except NotImplementedError as error:
160 self.handle_error(error)
161 continue
163 if machine_type.SLUG in machine_types: 163 ↛ 164line 163 didn't jump to line 164 because the condition on line 163 was never true
164 self.handle_error(
165 ValueError(
166 f"Cannot re-register machine type '{machine_type.SLUG}'"
167 )
168 )
169 continue
171 machine_types[machine_type.SLUG] = machine_type
172 base_drivers.append(machine_type.base_driver)
174 except Exception as error:
175 log_error(
176 'discover_machine_types',
177 plugin=plugin.slug,
178 scope='MachineRegistry',
179 )
180 self.handle_error(error)
182 self.machine_types = machine_types
183 self.base_drivers = base_drivers
185 logger.debug('Found %s machine types', len(self.machine_types.keys()))
187 def discover_drivers(self):
188 """Discovers all machine drivers by discovering all plugins which implement the Machine mixin class."""
189 from plugin import PluginMixinEnum
190 from plugin.registry import registry as plugin_registry
192 logger.debug('Collecting machine drivers')
193 drivers: dict[str, type[BaseDriver]] = {}
195 for plugin in plugin_registry.with_mixin(PluginMixinEnum.MACHINE):
196 try:
197 for driver in plugin.get_machine_drivers(): 197 ↛ 198line 197 didn't jump to line 198 because the loop on line 197 never started
198 if not issubclass(driver, BaseDriver):
199 logger.error(
200 'INVE-E12: Plugin %s returned invalid driver type',
201 plugin.slug,
202 )
203 continue
205 try:
206 driver.validate()
207 except NotImplementedError as error:
208 self.handle_error(error)
209 continue
211 if driver.SLUG in drivers:
212 self.handle_error(
213 ValueError(f"Cannot re-register driver '{driver.SLUG}'")
214 )
215 continue
217 drivers[driver.SLUG] = driver
218 except Exception as error:
219 log_error(
220 'discover_drivers', plugin=plugin.slug, scope='MachineRegistry'
221 )
222 self.handle_error(error)
224 self.drivers = drivers
226 logger.debug('Found %s machine drivers', len(self.drivers.keys()))
228 @machine_registry_entrypoint()
229 def get_driver_instance(self, slug: str):
230 """Return or create a driver instance if needed."""
231 if slug not in self.driver_instances:
232 driver = self.drivers.get(slug, None)
233 if driver is None:
234 return None
236 self.driver_instances[slug] = driver()
238 return self.driver_instances.get(slug, None)
240 @machine_registry_entrypoint()
241 def load_machines(self, main: bool = False):
242 """Load all machines defined in the database into the machine registry."""
243 # Imports need to be in this level to prevent early db model imports
245 try:
246 from machine.models import MachineConfig
248 for machine_config in MachineConfig.objects.all(): 248 ↛ 249line 248 didn't jump to line 249 because the loop on line 248 never started
249 self.add_machine(
250 machine_config, initialize=False, update_registry_hash=False
251 )
252 except (OperationalError, ProgrammingError):
253 logger.warning('Database is not ready - cannot load machines')
255 self._update_registry_hash()
256 return
258 # initialize machines only in main thread
259 if main:
260 # initialize drivers
261 for driver in self.driver_instances.values(): 261 ↛ 262line 261 didn't jump to line 262 because the loop on line 261 never started
262 driver.init_driver()
264 # initialize machines after all machine instances were created
265 for machine in self.machines.values(): 265 ↛ 266line 265 didn't jump to line 266 because the loop on line 265 never started
266 if machine.active:
267 machine.initialize()
269 logger.info('Initialized %s machines', len(self.machines.keys()))
270 else:
271 logger.info('Loaded %s machines', len(self.machines.keys()))
273 self._update_registry_hash()
275 def reload_machines(self, main: bool = False):
276 """Reload all machines from the database."""
277 self.drivers = {}
278 self.driver_instances = {}
279 self.machines = {}
281 InvenTree.cache.set_session_cache('machine_registry_checked', False)
283 self.set_shared_state('errors', [])
285 self.discover_machine_types()
286 self.discover_drivers()
287 self.load_machines(main=main)
289 @machine_registry_entrypoint()
290 def add_machine(self, machine_config, initialize=True, update_registry_hash=True):
291 """Add a machine to the machine registry."""
292 machine_type = self.machine_types.get(machine_config.machine_type, None)
293 if machine_type is None: 293 ↛ 297line 293 didn't jump to line 297 because the condition on line 293 was always true
294 self.handle_error(f"Machine type '{machine_config.machine_type}' not found")
295 return
297 machine: BaseMachineType = machine_type(machine_config)
298 self.machines[str(machine.pk)] = machine
300 if initialize and machine.active:
301 machine.initialize()
303 if update_registry_hash:
304 self._update_registry_hash()
306 @machine_registry_entrypoint()
307 def update_machine(
308 self, old_machine_state, machine_config, update_registry_hash=True
309 ):
310 """Notify the machine about an update."""
311 if machine := machine_config.machine:
312 machine.update(old_machine_state)
314 if update_registry_hash:
315 self._update_registry_hash()
317 @machine_registry_entrypoint()
318 def restart_machine(self, machine):
319 """Restart a machine."""
320 machine.restart()
322 @machine_registry_entrypoint()
323 def remove_machine(self, machine: BaseMachineType):
324 """Remove a machine from the registry."""
325 self.machines.pop(str(machine.pk), None)
326 self._update_registry_hash()
328 @machine_registry_entrypoint(default_value=False)
329 def get_machines(self, **kwargs):
330 """Get loaded machines from registry (By default only initialized machines).
332 Kwargs:
333 name: Machine name
334 machine_type: Machine type definition (class)
335 driver: Machine driver (class)
336 initialized (bool | None): use None to get all machines (default: True)
337 active: (bool)
338 base_driver: base driver (class)
339 """
340 allowed_fields = [
341 'name',
342 'machine_type',
343 'driver',
344 'initialized',
345 'active',
346 'base_driver',
347 ]
349 if 'initialized' not in kwargs:
350 kwargs['initialized'] = True
351 if kwargs['initialized'] is None:
352 del kwargs['initialized']
354 def filter_machine(machine: BaseMachineType):
355 for key, value in kwargs.items():
356 if key not in allowed_fields:
357 raise ValueError(
358 f"'{key}' is not a valid filter field for registry.get_machines."
359 )
361 # check if current driver is subclass from base_driver
362 if key == 'base_driver':
363 if machine.driver and not issubclass(
364 machine.driver.__class__, value
365 ):
366 return False
368 # check if current machine is subclass from machine_type
369 elif key == 'machine_type':
370 if issubclass(machine.__class__, value):
371 return False
373 # check attributes of machine
374 elif value != getattr(machine, key, None):
375 return False
377 return True
379 return list(filter(filter_machine, self.machines.values()))
381 @machine_registry_entrypoint(default_value=[])
382 def get_machine_types(self):
383 """Get all machine types."""
384 return list(self.machine_types.values())
386 @machine_registry_entrypoint()
387 def get_machine(self, pk: str | UUID) -> Optional[BaseMachineType]:
388 """Get machine from registry by pk."""
389 return self.machines.get(str(pk), None)
391 @machine_registry_entrypoint(default_value=[])
392 def get_driver_types(self, machine_type: Optional[str] = None):
393 """Return a list of all registered driver types.
395 Arguments:
396 machine_type: Optional machine type to filter drivers by their machine type
397 """
398 return [
399 driver
400 for driver in self.drivers.values()
401 if machine_type is None or driver.machine_type == machine_type
402 ]
404 @machine_registry_entrypoint(default_value=[])
405 def get_drivers(self, machine_type: Optional[str] = None):
406 """Get all drivers for a specific machine type.
408 Arguments:
409 machine_type: Optional machine type to filter drivers by their machine type
410 """
411 return [
412 driver
413 for driver in self.driver_instances.values()
414 if machine_type is None or driver.machine_type == machine_type
415 ]
417 def _calculate_registry_hash(self):
418 """Calculate a hash of the machine registry state."""
419 from hashlib import md5
421 from plugin import registry as plugin_registry
423 data = md5()
425 # If the plugin registry has changed, the machine registry hash will change
426 plugin_registry.update_plugin_hash()
427 current_hash = plugin_registry.registry_hash
428 if current_hash: 428 ↛ 431line 428 didn't jump to line 431 because the condition on line 428 was always true
429 data.update(current_hash.encode())
431 for pk, machine in self.machines.items(): 431 ↛ 432line 431 didn't jump to line 432 because the loop on line 431 never started
432 data.update(str(pk).encode())
433 try:
434 data.update(str(machine.machine_config.active).encode())
435 except:
436 # machine does not exist anymore, hash will be different
437 pass
439 return str(data.hexdigest())
441 def _check_reload(self):
442 """Check if the registry needs to be reloaded, and reload it."""
443 from plugin import registry as plg_registry
445 do_reload: bool = False
446 plugin_registry_hash = getattr(self, '_plugin_registry_hash', None)
448 if plugin_registry_hash != plg_registry.registry_hash:
449 do_reload = True
451 if not self._hash:
452 self._hash = self._calculate_registry_hash()
454 try:
455 reg_hash = get_global_setting('_MACHINE_REGISTRY_HASH', '', create=False)
456 except Exception as exc:
457 logger.exception('Failed to get machine registry hash: %s', exc)
458 return False
460 if reg_hash and reg_hash != self._hash: 460 ↛ 461line 460 didn't jump to line 461 because the condition on line 460 was never true
461 logger.info('Machine registry has changed - reloading machines')
462 do_reload = True
464 if do_reload:
465 self.reload_machines()
467 return do_reload
469 def _update_registry_hash(self):
470 """Save the current registry hash."""
471 from plugin import registry as plg_registry
473 self._hash = self._calculate_registry_hash()
474 self._plugin_registry_hash = plg_registry.registry_hash
476 try:
477 old_hash = get_global_setting('_MACHINE_REGISTRY_HASH')
478 except Exception:
479 old_hash = None
481 if old_hash != self._hash:
482 try:
483 logger.info('Updating machine registry hash: %s', self._hash)
484 set_global_setting('_MACHINE_REGISTRY_HASH', self._hash)
485 except (IntegrityError, OperationalError, ProgrammingError):
486 pass
487 except Exception as exc:
488 logger.exception('Failed to update machine registry hash: %s', exc)
490 @machine_registry_entrypoint()
491 def call_machine_function(
492 self, machine_id: str, function_name: str, *args, **kwargs
493 ) -> Any:
494 """Call a named function against a machine instance.
496 Arguments:
497 machine_id: The UUID of the machine to call the function against
498 function_name: The name of the function to call
499 """
500 logger.info('call_machine_function: %s -> %s', machine_id, function_name)
502 raise_error = kwargs.pop('raise_error', True)
504 # Fetch the machine instance based on the provided UUID
505 machine = self.get_machine(machine_id)
507 if not machine:
508 if raise_error:
509 raise AttributeError(f"Machine '{machine_id}' not found")
510 return
512 # Fetch the driver instance based on the machine driver
513 driver = machine.driver
515 if not driver:
516 if raise_error:
517 raise AttributeError(f"Machine '{machine_id}' has no specified driver")
518 return
520 # The function must be registered against the driver
521 func = getattr(driver, function_name)
523 if not func or not callable(func):
524 if raise_error:
525 raise AttributeError(
526 f"Driver '{driver.SLUG}' has no callable method '{function_name}'"
527 )
528 return
530 return func(machine, *args, **kwargs)
533registry: MachineRegistry = MachineRegistry()
536def call_machine_function(machine_id: str, function: str, *args, **kwargs):
537 """Global helper function to call a specific function on a machine instance.
539 Arguments:
540 machine_id: The UUID of the machine to call the function against
541 function: The name of the function to call
542 *args: Positional arguments to pass to the function
543 **kwargs: Keyword arguments to pass to the function
544 """
545 return registry.call_machine_function(machine_id, function, *args, **kwargs)