cumulusci-plus 5.0.0__py3-none-any.whl

cumulusci/tasks/bulkdata/snowfakery_utils/queue_manager.py

@@ -0,0 +1,478 @@
+import queue
+import random
+import shutil
+import time
+import typing as T
+from collections import defaultdict
+from multiprocessing import Lock
+from pathlib import Path
+
+import cumulusci.core.exceptions as exc
+from cumulusci.core.config import OrgConfig
+from cumulusci.tasks.bulkdata.generate_from_yaml import GenerateDataFromYaml
+from cumulusci.tasks.bulkdata.load import LoadData
+from cumulusci.utils.parallel.task_worker_queues.parallel_worker_queue import (
+    WorkerQueue,
+    WorkerQueueConfig,
+)
+
+from .snowfakery_run_until import PortionGenerator
+from .snowfakery_working_directory import SnowfakeryWorkingDirectory
+from .subtask_configurator import SubtaskConfigurator
+
+# number of portions we will allow to be on-disk waiting to be loaded
+# higher numbers use more disk space.
+LOAD_QUEUE_SIZE = 15
+# more loader workers than generators because they spend so much time
+# waiting for responses. 4:1 is an experimentally derived ratio
+WORKER_TO_LOADER_RATIO = 4
+
+
+class SnowfakeryChannelManager:
+    """The channels describe the 'shape' of the communication
+
+    The normal case is a single, parallelized, bulk channel,
+    multi-threaded on client and server, using a single user
+    account.
+
+    Using .load.yml you can add more channels, utilizing
+    more user accounts which can speed up throughput in
+    a few cases."""
+
+    def __init__(
+        self,
+        subtask_configurator,
+        *,
+        project_config,
+        logger,
+    ):
+        # Look at the docstring on get_results_report to understand
+        # what this queue is for.
+        #
+        # Be careful to use a Queue class appropriate to
+        # the spawn type (thread, process) you're using.
+        #
+        # Snowfakery runs its loader in threads, so queue.Queue()
+        # works.
+        #
+        # multiprocessing.Manager().Queue() also seems to work,
+        # and work across processes, (PR #3080) but it's dramatically
+        # slower. See attachment to PR #3076
+        self.results_reporter = queue.Queue()
+        self.channels = []
+        self.project_config = project_config
+        self.logger = logger
+        self.subtask_configurator = subtask_configurator
+        self.start_time = time.time()
+
+    def add_channel(
+        self,
+        org_config: OrgConfig,
+        num_generator_workers: int,
+        num_loader_workers: T.Optional[int],
+        working_directory: Path,
+        recipe_options: dict,
+    ):
+        if not num_loader_workers:
+            num_loader_workers = num_generator_workers * WORKER_TO_LOADER_RATIO
+
+        self.channels.append(
+            Channel(
+                project_config=self.project_config,
+                org_config=org_config,
+                num_generator_workers=num_generator_workers,
+                num_loader_workers=num_loader_workers,
+                working_directory=working_directory,
+                subtask_configurator=self.subtask_configurator,
+                logger=self.logger,
+                results_reporter=self.results_reporter,
+                recipe_options=recipe_options,
+            )
+        )
+
+    def tick(
+        self,
+        upload_status,
+        template_path: Path,
+        tempdir: Path,
+        portions: PortionGenerator,
+        get_upload_status: T.Callable,
+    ):
+        """Called every few seconds, to make new data generators if needed."""
+        # advance each channel's queues
+        for channel in self.channels:
+            channel.tick()
+
+        # populate channels that have space
+        self.assign_work_to_channels(
+            upload_status,
+            template_path,
+            tempdir,
+            portions,
+            get_upload_status,
+        )
+
+    def assign_work_to_channels(
+        self,
+        upload_status,
+        template_path: Path,
+        tempdir: Path,
+        portions: PortionGenerator,
+        get_upload_status: T.Callable,
+    ):
+        def channel_with_free_space():
+            return [channel for channel in self.channels if not channel.full]
+
+        new_workers = [True]  # initial value to get into the while loop
+        channels = channel_with_free_space()
+        while (
+            channels
+            and any(new_workers)
+            and not portions.done(upload_status.total_sets_working_on_or_uploaded)
+        ):
+            # randomize channels so the first one doesn't always grab all slots
+            random.shuffle(channels)
+            new_workers = [
+                channel.make_new_worker(
+                    template_path, tempdir, portions, get_upload_status
+                )
+                for channel in channels
+            ]
+            channels = channel_with_free_space()
+
+    def get_upload_status(self, batch_size, sets_finished_while_generating_template):
+        """Combine information from the different data sources into a single "report".
+
+        Useful for debugging but also for making decisions about what to do next."""
+        summed_statuses = defaultdict(int)
+        for channel in self.channels:
+            for key, value in channel.get_upload_status_for_channel().items():
+                summed_statuses[key] += value
+
+        summed_statuses["sets_finished"] += sets_finished_while_generating_template
+        return UploadStatus(
+            target_count=self.subtask_configurator.run_until.gap,
+            elapsed_seconds=int(self.elapsed_seconds()),
+            batch_size=batch_size,
+            channels=len(self.channels),
+            **summed_statuses,
+        )
+
+    def elapsed_seconds(self):
+        return time.time() - self.start_time
+
+    def failure_descriptions(self) -> T.List[str]:
+        """Log failures from sub-processes to main process"""
+        ret = []
+        for channel in self.channels:
+            ret.extend(channel.failure_descriptions())
+        return ret
+
+    def check_finished(self) -> bool:
+        for channel in self.channels:
+            channel.tick()
+        return all([channel.check_finished() for channel in self.channels])
+
+    def get_results_report(self, block=False):
+        """
+        This is a realtime reporting channel which could, in theory, be updated
+        before sub-tasks finish. Currently no sub-tasks are coded to do that.
+
+        The logical next step is to allow LoadData to monitor steps one by
+        one or even batches one by one.
+
+        Note that until we implement that, we are paying the complexity
+        cost of a real-time channel but not getting the benefits of it."""
+        return self.results_reporter.get(block=block)
+
+
+class Channel:
+    def __init__(
+        self,
+        *,
+        project_config,
+        org_config,
+        num_generator_workers: int,
+        num_loader_workers: int,
+        working_directory: Path,
+        subtask_configurator: SubtaskConfigurator,
+        logger,
+        recipe_options=None,
+        results_reporter=None,
+    ):
+        self.project_config = project_config
+        self.org_config = org_config
+        self.num_generator_workers = num_generator_workers
+        self.num_loader_workers = num_loader_workers
+        self.working_directory = working_directory
+        self.subtask_configurator = subtask_configurator
+        self.run_until = subtask_configurator.run_until
+        self.logger = logger
+        self.results_reporter = results_reporter
+        self.filesystem_lock = Lock()
+        self.job_counter = 0
+        recipe_options = recipe_options or {}
+        self._configure_queues(recipe_options)
+
+    def _configure_queues(self, recipe_options):
+        """Configure two ParallelWorkerQueues for datagen and dataload"""
+        try:
+            connected_app = self.project_config.keychain.get_service(
+                "connected_app", self.org_config.connected_app
+            )
+        except exc.ServiceNotConfigured:  # pragma: no cover
+            # to discuss...when can this happen? What are the consequences?
+            connected_app = None
+
+        def data_generator_opts_callback(*args, **kwargs):
+            return self.subtask_configurator.data_generator_opts(
+                *args, recipe_options=recipe_options, **kwargs
+            )
+
+        data_gen_q_config = WorkerQueueConfig(
+            project_config=self.project_config,
+            org_config=self.org_config,
+            connected_app=connected_app,
+            redirect_logging=True,
+            # processes are better for compute-heavy tasks (in Python)
+            spawn_class=WorkerQueue.Process,
+            parent_dir=self.working_directory,
+            name="data_gen",
+            task_class=GenerateDataFromYaml,
+            make_task_options=data_generator_opts_callback,
+            queue_size=0,
+            num_workers=self.num_generator_workers,
+        )
+        # datagen queues do not get a result reporter because
+        # a) we are less curious about how many records have
+        #    been generated than we aare about how many are loaded
+        # b) finding a queue type which is not prone to race conditions
+        #    or perf slowdowns with "Process" task types (sub-processes)
+        # is a challenge.
+        self.data_gen_q = WorkerQueue(data_gen_q_config, self.filesystem_lock)
+
+        load_data_q_config = WorkerQueueConfig(
+            project_config=self.project_config,
+            org_config=self.org_config,
+            connected_app=connected_app,
+            redirect_logging=True,
+            spawn_class=WorkerQueue.Thread,
+            parent_dir=self.working_directory,
+            name="data_load",
+            task_class=LoadData,
+            make_task_options=self.subtask_configurator.data_loader_opts,
+            queue_size=LOAD_QUEUE_SIZE,
+            num_workers=self.num_loader_workers,
+            rename_directory=self.data_loader_new_directory_name,
+        )
+        self.load_data_q = WorkerQueue(
+            load_data_q_config, self.filesystem_lock, self.results_reporter
+        )
+
+        self.data_gen_q.feeds_data_to(self.load_data_q)
+        return self.data_gen_q, self.load_data_q
+
+    def data_loader_new_directory_name(self, working_directory):
+        return data_loader_new_directory_name(working_directory, self.run_until)
+
+    def tick(
+        self,
+    ):
+        """Called every few seconds, to make new data generators if needed."""
+        self.data_gen_q.tick()
+
+    @property
+    def full(self):
+        return self.data_gen_q.full
+
+    def make_new_worker(
+        self,
+        template_path: Path,
+        tempdir: Path,
+        portions: PortionGenerator,
+        get_upload_status: T.Callable,
+    ):
+        if (
+            self.data_gen_q.num_free_workers and self.data_gen_q.full
+        ):  # pragma: no cover
+            # TODO: investigate the consequences of taking this branch out
+            self.logger.info("Waiting before datagen (load queue is full)")
+        else:
+            upload_status = get_upload_status(
+                portions.next_batch_size,
+            )
+            self.job_counter += 1
+            batch_size = portions.next_batch(
+                upload_status.total_sets_working_on_or_uploaded
+            )
+            if not batch_size:
+                self.logger.info("All scheduled portions generated and being uploaded")
+                return
+
+            job_dir = self.generator_data_dir(
+                self.job_counter, template_path, batch_size, tempdir
+            )
+            self.data_gen_q.push(job_dir)
+            return job_dir
+
+    def generator_data_dir(self, idx, template_path, batch_size, parent_dir):
+        """Create a new generator directory with a name based on index and batch_size"""
+        assert batch_size > 0
+        data_dir = parent_dir / (str(idx) + "_" + str(batch_size))
+        shutil.copytree(template_path, data_dir)
+        return data_dir
+
+    def get_upload_status_for_channel(self):
+        with self.filesystem_lock:
+
+            def set_count_from_names(names):
+                return sum(int(name.split("_")[1]) for name in names)
+
+            return {
+                "sets_queued_to_be_generated": set_count_from_names(
+                    self.data_gen_q.queued_jobs
+                ),
+                "sets_being_generated": set_count_from_names(
+                    self.data_gen_q.inprogress_jobs
+                ),
+                "sets_queued_for_loading": set_count_from_names(
+                    self.load_data_q.queued_jobs
+                ),
+                # note that these may count as already imported in the org
+                "sets_being_loaded": set_count_from_names(
+                    self.load_data_q.inprogress_jobs
+                ),
+                "max_num_loader_workers": self.num_loader_workers,
+                "max_num_generator_workers": self.num_generator_workers,
+                # todo: use row-level result from org load for better accuracy
+                "sets_finished": set_count_from_names(self.load_data_q.outbox_jobs),
+                "sets_failed": len(self.load_data_q.failed_jobs),
+                # TODO: are these two redundant?
+                "inprogress_generator_jobs": len(self.data_gen_q.inprogress_jobs),
+                "inprogress_loader_jobs": len(self.load_data_q.inprogress_jobs),
+                "data_gen_free_workers": self.data_gen_q.num_free_workers,
+            }
+
+    def failure_descriptions(self) -> T.List[str]:
+        """Log failures from sub-processes to main process"""
+        failure_dirs = set(
+            self.load_data_q.failed_job_dirs + self.data_gen_q.failed_job_dirs
+        )
+
+        def error_from_dir(failure_dir: Path) -> T.Optional[str]:
+            f = Path(failure_dir) / "exception.txt"
+            if not f.exists():  # pragma: no cover
+                return
+            return f.read_text(encoding="utf-8").strip().split("\n")[-1]
+
+        errors = [error_from_dir(failure_dir) for failure_dir in failure_dirs]
+        return [error for error in errors if error is not None]
+
+    def check_finished(self) -> bool:
+        self.data_gen_q.tick()
+        with self.filesystem_lock:
+            still_running = (
+                len(
+                    self.data_gen_q.workers
+                    + self.data_gen_q.queued_job_dirs
+                    + self.data_gen_q.inprogress_jobs
+                    + self.load_data_q.workers
+                    + self.load_data_q.inprogress_jobs
+                    + self.load_data_q.queued_job_dirs
+                )
+                > 0
+            )
+        return not still_running
+
+
+# TODO: This function is actually based on the number generated,
+#       because it is called before the load.
+#       If there are row errors, it will drift out of correctness
+#       Code needs to be updated to rename again after load.
+#       Or move away from using directory names for math altogether.
+def data_loader_new_directory_name(working_dir: Path, run_until):
+    """Change the directory name to reflect the true number of sets created."""
+
+    if not run_until.sobject_name:
+        return working_dir
+
+    wd = SnowfakeryWorkingDirectory(working_dir)
+    count = wd.get_record_counts()[run_until.sobject_name]
+
+    path, _ = str(working_dir).rsplit("_", 1)
+    new_working_dir = Path(path + "_" + str(count))
+    return new_working_dir
+
+
+class UploadStatus(T.NamedTuple):
+    """Single "report" of the current status of our processes."""
+
+    batch_size: int
+    sets_being_generated: int
+    sets_queued_to_be_generated: int
+    sets_being_loaded: int
+    sets_queued_for_loading: int
+    sets_finished: int
+    target_count: int
+    max_num_loader_workers: int
+    max_num_generator_workers: int
+    elapsed_seconds: int
+    sets_failed: int
+    inprogress_generator_jobs: int
+    inprogress_loader_jobs: int
+    data_gen_free_workers: int
+    channels: int
+
+    @property
+    def total_in_flight(self):
+        return (
+            self.sets_queued_to_be_generated
+            + self.sets_being_generated
+            + self.sets_queued_for_loading
+            + self.sets_being_loaded
+        )
+
+    @property
+    def total_sets_working_on_or_uploaded(
+        self,
+    ):
+        return self.total_in_flight + self.sets_finished
+
+    def _display(self, detailed=False):
+        most_important_stats = [
+            "target_count",
+            "total_sets_working_on_or_uploaded",
+            "sets_finished",
+            "sets_failed",
+        ]
+        if self.channels > 1:
+            most_important_stats.append("channels")
+
+        queue_stats = [
+            "inprogress_generator_jobs",
+            "inprogress_loader_jobs",
+        ]
+
+        def display_stats(keys):
+            def format(a):
+                return f"{a.replace('_', ' ').title()}: {getattr(self, a):,}"
+
+            return (
+                "\n"
+                + "\n".join(
+                    format(a)
+                    for a in keys
+                    if not a[0] == "_" and not callable(getattr(self, a))
+                )
+                + "\n"
+            )
+
+        rc = display_stats(most_important_stats)
+        rc += "\n   ** Workers **\n"
+        rc += display_stats(queue_stats)
+        if detailed:
+            rc += "\n   ** Internals **\n"
+            rc += display_stats(
+                set(dir(self)) - (set(most_important_stats) & set(queue_stats))
+            )
+        return rc

cumulusci/tasks/bulkdata/snowfakery_utils/snowfakery_run_until.py

@@ -0,0 +1,141 @@
+import typing as T
+
+from cumulusci.core import exceptions as exc
+
+
+class RunUntilBase:
+    # subclasses need to fill in these two fields
+    sobject_name: str = None
+    target: int = None
+    gap: int = None
+
+    def set_target_and_gap(
+        self, sobject_name: T.Optional[str], num_as_str: T.Union[str, int]
+    ):
+        self.sobject_name = sobject_name
+        try:
+            self.target = self.gap = int(num_as_str)
+        except (TypeError, ValueError):
+            raise exc.TaskOptionsError(f"{num_as_str} is not a number")
+
+    def split_pair(self, param):
+        parts = param.split(":")
+        if len(parts) != 2:
+            raise exc.TaskOptionsError(
+                f"{param} is in the wrong format for {self.option_name}"
+            )
+        return parts
+
+    @property
+    def nothing_to_do(self):
+        return self.gap <= 0
+
+
+class RunUntilRecipeRepeated(RunUntilBase):
+    option_name = "--run-until-recipe-repeated"
+    nothing_to_do_because = "You asked for zero repetitions."
+
+    def __init__(self, sf, param):
+        self.set_target_and_gap(None, param)
+
+
+class RunUntilRecordsLoaded(RunUntilBase):
+    option_name = "--run-until-records-loaded"
+    nothing_to_do_because = "You asked for zero records."
+
+    def __init__(self, sf, param):
+        parts = self.split_pair(param)
+        self.set_target_and_gap(*parts)
+
+
+class RunUntilRecordInOrg(RunUntilRecordsLoaded):
+    option_name = "--run-until-records-in-org"
+
+    def __init__(self, sf, param):
+        parts = self.split_pair(param)
+        self.set_target_and_gap(*parts)
+        query = f"select count(Id) from {self.sobject_name}"
+        # TODO: fall back to other APIs
+        self.in_org_count = sf.query(query)["records"][0]["expr0"]
+        gap = self.target - int(self.in_org_count)
+        self.gap = max(gap, 0)
+
+    @property
+    def nothing_to_do_because(self):
+        return f"The org has {self.in_org_count} {self.sobject_name} records. You asked for {self.target}"
+
+
+COUNT_STRATEGIES = {
+    "run_until_recipe_repeated": RunUntilRecipeRepeated,
+    "run_until_records_loaded": RunUntilRecordsLoaded,
+    "run_until_records_in_org": RunUntilRecordInOrg,
+}
+
+
+def determine_run_until(options, sf):
+    selected_strategies = [
+        (strategy, options.get(strategy))
+        for strategy in COUNT_STRATEGIES.keys()
+        if options.get(strategy)
+    ]
+
+    if len(selected_strategies) > 1:
+        raise exc.TaskOptionsError(
+            "Please select only one of " + ", ".join(COUNT_STRATEGIES.keys())
+        )
+    elif not selected_strategies:
+        strategy_choice = ("run_until_recipe_repeated", "1")
+    else:
+        strategy_choice = selected_strategies[0]
+
+    strategy_name, param = strategy_choice
+    strategy_impl = COUNT_STRATEGIES[strategy_name]
+
+    return strategy_impl(sf, param)
+
+
+class PortionGenerator:
+    """
+    Logic to generate batch sizes and know when we are done.
+
+    Calling class must manage sets_created_so_far because it may
+    have information unknown to this class, e.g. failures,
+    processes running on other nodes, etc.
+    """
+
+    def __init__(
+        self,
+        target: int,
+        min_batch_size: int,
+        max_batch_size: int,
+        start_count: int = 0,
+    ):
+        self.target = target
+        self.min_batch_size = min_batch_size
+        self.max_batch_size = max_batch_size
+        self.next_batch_size = min_batch_size
+
+    def done(self, sets_created_so_far):
+        return self.gap(sets_created_so_far) <= 0
+
+    def gap(self, sets_created_so_far):
+        return self.target - sets_created_so_far
+
+    def next_batch(self, sets_created_so_far: int):
+        """Batch size starts at min_batch_size and grows toward max_batch_size
+        unless the count gets there first
+
+        The reason we grow is to try to engage the org's loader queue earlier
+        than we would if we waited for the first big batch to be generated.
+        """
+        self.sets_created = sets_created_so_far
+        # don't generate more records than needed to reach our target
+        # or current batch size
+        # or the max batch size
+        self.batch_size = min(
+            self.gap(sets_created_so_far), self.next_batch_size, self.max_batch_size
+        )
+        # don't generate fewer than zero
+        self.batch_size = max(self.batch_size, 0)
+        self.next_batch_size = int(self.next_batch_size * 1.1)
+        return self.batch_size

cumulusci/tasks/bulkdata/snowfakery_utils/snowfakery_working_directory.py

@@ -0,0 +1,53 @@
+import yaml
+from sqlalchemy import MetaData, Table, create_engine, func, select
+
+
+class SnowfakeryWorkingDirectory:
+    """Helper functions based on well-known filenames in CCI/Snowfakery working directories."""
+
+    def __init__(self, working_dir):
+        self.path = working_dir
+        self.mapping_file = working_dir / "temp_mapping.yml"
+        self.database_file = working_dir / "generated_data.db"
+        assert self.mapping_file.exists(), self.mapping_file
+        assert self.database_file.exists(), self.database_file
+        self.database_url = f"sqlite:///{self.database_file}"
+        self.continuation_file = f"{working_dir}/continuation.yml"
+
+    def setup_engine(self):
+        """Set up the database engine"""
+        engine = create_engine(self.database_url)
+
+        metadata = MetaData(engine)
+        metadata.reflect()
+        return engine, metadata
+
+    @property
+    def index(self) -> str:
+        return self.path.name.rsplit("_")[0]
+
+    def get_record_counts(self):
+        """Get record counts generated for this portion."""
+        engine, metadata = self.setup_engine()
+
+        with engine.connect() as connection:
+            record_counts = {
+                table_name: self._record_count_from_db(connection, table)
+                for table_name, table in metadata.tables.items()
+                if table_name[-6:] != "sf_ids"
+            }
+        # note that the template has its contents deleted so if the cache
+        # is ever removed, it will start to return {}
+        assert record_counts
+        return record_counts
+
+    def _record_count_from_db(self, connection, table: Table):
+        """Count rows in a table"""
+        stmt = select(func.count()).select_from(table)
+        result = connection.execute(stmt)
+        return next(result)[0]
+
+    def relevant_sobjects(self):
+        with open(self.mapping_file, encoding="utf-8") as f:
+            mapping = yaml.safe_load(f)
+            return [m.get("sf_object") for m in mapping.values() if m.get("sf_object")]