|
5 | 5 |
|
6 | 6 | import datetime |
7 | 7 | import logging |
| 8 | +import os |
8 | 9 | from typing import Any, Dict, List, Optional, Tuple, Union |
9 | 10 | from uuid import UUID |
10 | 11 |
|
@@ -437,8 +438,13 @@ def cleanup_nodes(self) -> bool: |
437 | 438 |
|
438 | 439 | # Perform operations until they fail due to scaleset getting locked |
439 | 440 | try: |
440 | | - self.reimage_nodes(to_reimage, NodeDisaposalStrategy.scale_in) |
441 | | - self.delete_nodes(to_delete, NodeDisaposalStrategy.scale_in) |
| 441 | + strategy_str = os.getenv("ONEFUZZ_NODE_DISPOSAL_STRATEGY", "scale_in") |
| 442 | + if strategy_str == "decomission": |
| 443 | + strategy = NodeDisaposalStrategy.decomission |
| 444 | + else: |
| 445 | + strategy = NodeDisaposalStrategy.scale_in |
| 446 | + self.reimage_nodes(to_reimage, strategy) |
| 447 | + self.delete_nodes(to_delete, strategy) |
442 | 448 | except UnableToUpdate: |
443 | 449 | logging.info( |
444 | 450 | SCALESET_LOG_PREFIX |
@@ -598,17 +604,23 @@ def delete_nodes( |
598 | 604 | else: |
599 | 605 | machine_ids.add(node.machine_id) |
600 | 606 |
|
601 | | - logging.info( |
602 | | - SCALESET_LOG_PREFIX + "deleting nodes scaleset_id:%s machine_id:%s", |
603 | | - self.scaleset_id, |
604 | | - machine_ids, |
605 | | - ) |
606 | | - delete_vmss_nodes(self.scaleset_id, machine_ids) |
607 | | - for node in nodes: |
608 | | - if node.machine_id in machine_ids: |
609 | | - node.delete() |
610 | | - if disposal_strategy == NodeDisaposalStrategy.scale_in: |
| 607 | + if disposal_strategy == NodeDisaposalStrategy.decomission: |
| 608 | + logging.info(SCALESET_LOG_PREFIX + "decomissioning nodes") |
| 609 | + for node in nodes: |
| 610 | + if node.machine_id in machine_ids: |
611 | 611 | node.release_scale_in_protection() |
| 612 | + else: |
| 613 | + logging.info( |
| 614 | + SCALESET_LOG_PREFIX + "deleting nodes scaleset_id:%s machine_id:%s", |
| 615 | + self.scaleset_id, |
| 616 | + machine_ids, |
| 617 | + ) |
| 618 | + delete_vmss_nodes(self.scaleset_id, machine_ids) |
| 619 | + for node in nodes: |
| 620 | + if node.machine_id in machine_ids: |
| 621 | + node.delete() |
| 622 | + if disposal_strategy == NodeDisaposalStrategy.scale_in: |
| 623 | + node.release_scale_in_protection() |
612 | 624 |
|
613 | 625 | def reimage_nodes( |
614 | 626 | self, nodes: List[Node], disposal_strategy: NodeDisaposalStrategy |
@@ -659,18 +671,24 @@ def reimage_nodes( |
659 | 671 | ) |
660 | 672 | return |
661 | 673 |
|
662 | | - result = reimage_vmss_nodes(self.scaleset_id, machine_ids) |
663 | | - if isinstance(result, Error): |
664 | | - raise Exception( |
665 | | - "unable to reimage nodes: %s:%s - %s" |
666 | | - % (self.scaleset_id, machine_ids, result) |
667 | | - ) |
668 | | - |
669 | | - for node in nodes: |
670 | | - if node.machine_id in machine_ids: |
671 | | - node.delete() |
672 | | - if disposal_strategy == NodeDisaposalStrategy.scale_in: |
| 674 | + if disposal_strategy == NodeDisaposalStrategy.decomission: |
| 675 | + logging.info(SCALESET_LOG_PREFIX + "decomissioning nodes") |
| 676 | + for node in nodes: |
| 677 | + if node.machine_id in machine_ids: |
673 | 678 | node.release_scale_in_protection() |
| 679 | + else: |
| 680 | + result = reimage_vmss_nodes(self.scaleset_id, machine_ids) |
| 681 | + if isinstance(result, Error): |
| 682 | + raise Exception( |
| 683 | + "unable to reimage nodes: %s:%s - %s" |
| 684 | + % (self.scaleset_id, machine_ids, result) |
| 685 | + ) |
| 686 | + |
| 687 | + for node in nodes: |
| 688 | + if node.machine_id in machine_ids: |
| 689 | + node.delete() |
| 690 | + if disposal_strategy == NodeDisaposalStrategy.scale_in: |
| 691 | + node.release_scale_in_protection() |
674 | 692 |
|
675 | 693 | def set_shutdown(self, now: bool) -> None: |
676 | 694 | if now: |
|
0 commit comments