text: """
  Execute a partition strategy using multiple Gretel jobs.
  
  This module will consume a dataset and create a partitioning strategy. This strategy will be saved
  to a JSON file. Within each partition there is a ``ctx`` object that will contain metadta about the
  various models and handlers used. The context object has the following shape:
  
  {
      "artifact": "artifact used to train model",
      "status": "one of Gretel's job statuses",
      "model_id": "the model ID",
      "attempt": "number of times the model has attempted training",
      "sqs": "an object that contains SQS information from Gretel's API",
      "last_handler": {
          "artifact": "any seed data that was used for generation",
          "attempt": "number of tries to generate",
          "status": "a Gretel job status",
          "handler_id": "the record handler ID",
          "num_records": "how many records were generated, automatically overloaded if seeds are provided"
      }
  }
  """
  from __future__ import annotations
  
  import json
  import logging
  import tempfile
  import time
  from collections import Counter
  from concurrent.futures import ALL_COMPLETED, ThreadPoolExecutor, wait
  from copy import deepcopy
  from dataclasses import dataclass
  from functools import wraps
  from pathlib import Path
  from typing import List, Optional, Union
  
  import pandas as pd
  import smart_open
  from gretel_client.projects import Project
  from gretel_client.projects.jobs import ACTIVE_STATES
  from gretel_client.projects.models import Model, Status
  from gretel_client.projects.records import RecordHandler
  from gretel_client.rest import ApiException
  from gretel_client.users.users import get_me
  
  from gretel_trainer.strategy import (Partition, PartitionConstraints,
                                       PartitionStrategy)
  
  MODEL_ID = "model_id"
  HANDLER_ID = "handler_id"
  STATUS = "status"
  ARTIFACT = "artifact"
  SQS = "sqs"
  ATTEMPT = "attempt"
  HANDLER = "last_handler"
  NUM_RECS = "num_records"
  
  logger = logging.getLogger(__name__)
  logger.setLevel("DEBUG")
  
  
  @dataclass
  class ArtifactResult:
      id: str
      record_count: int
  
  
  @dataclass
  class GenPayload:
      num_records: int
      seed_df: Optional[pd.DataFrame] = None
      seed_artifact_id: Optional[str] = None
      max_invalid: Optional[int] = None
  
  
  def _needs_load(func):
      @wraps(func)
      def wrapper(inst: StrategyRunner, *args, **kwargs):
          if not inst._loaded:
              inst.load()
          return func(inst, *args, **kwargs)
  
      return wrapper
  
  
  @dataclass
  class RemoteDFPayload:
      partition: int
      slot: int
      job_type: str
      uid: str
      handler_uid: str
      project: Project
      artifact_type: str
      df: pd.DataFrame = None
  
  
  def _remote_dataframe_fetcher(payload: RemoteDFPayload) -> RemoteDFPayload:
      # We need the model object no matter what
      model = Model(payload.project, model_id=payload.uid)
      job = model
  
      # if we are downloading handler data, we reset our job
      # to the specific handler object
      if payload.job_type == "run":
          job = RecordHandler(model, record_id=payload.handler_uid)
  
      download_url = job.get_artifact_link(payload.artifact_type)
      payload.df = pd.read_csv(download_url, compression="gzip")
      return payload
  
  
  def _maybe_submit_job(
      job: Union[Model, RecordHandler]
  ) -> Optional[Union[Model, RecordHandler]]:
      try:
          job = job.submit_cloud()
      except ApiException as err:
          if "Maximum number of" in str(err):
              logger.warning(
                  "Rate limiting: Max jobs created, skipping new job for now..."
              )
          return None
  
      return job
  
  
  class StrategyRunner:
  
      _df: pd.DataFrame
      _cache_file: Path
      _constraints = PartitionConstraints
      _strategy = PartitionStrategy
      _model_config: dict
      _max_jobs_active: int
      _project: Project
      _loaded: bool
      _artifacts: List[str]
      _cache_overwrite: bool
      _max_artifacts: int = 25
      _status_counter: Counter
      _handler_status_counter: Counter
      _error_retry_limit: int
      strategy_id: str
  
      def __init__(
          self,
          *,
          strategy_id: str,
          df: pd.DataFrame,
          cache_file: Union[str, Path],
          cache_overwrite: bool = False,
          model_config: dict,
          partition_constraints: PartitionConstraints,
          project: Project,
          error_retry_limit: int = 3,
      ):
          self._df = df
          self._cache_file = Path(cache_file)
          self._constraints = partition_constraints
          self._model_config = model_config
          self._project = project
          self._loaded = False
          self._cache_overwrite = cache_overwrite
          self._artifacts = []
          self.strategy_id = strategy_id
          self._status_counter = Counter()
          self._error_retry_limit = error_retry_limit
  
      def load(self):
          """Hydrate the instance before we can start
          doing work, must be called after init
          """
          if self._loaded:
              return
  
          self._refresh_max_job_capacity()
  
          # If the cache file exists, we'll try and load an existing
          # strategy. If not, we'll create a new strategy with the
          # provided constraints.
  
          if self._cache_file.exists() and not self._cache_overwrite:
              self._strategy = PartitionStrategy.from_disk(self._cache_file)
          else:
              self._strategy = PartitionStrategy.from_dataframe(
                  self.strategy_id, self._df, self._constraints
              )
  
          self._strategy.save_to(self._cache_file, overwrite=True)
  
          self._loaded = True
  
      @classmethod
      def from_completed(
          cls, project: Project, cache_file: Union[str, Path]
      ) -> StrategyRunner:
          cache_file = Path(cache_file)
          if not cache_file.exists():
              raise ValueError("cache file does not exist")
  
          inst = cls(
              strategy_id="__none__",
              df=None,
              cache_file=cache_file,
              model_config=None,
              partition_constraints=None,
              project=project,
          )
          inst.load()
          return inst
  
      def _update_job_status(self):
          # Get all jobs that have been created, we can do this
          # by just searching for any partitions have have a "model_id"
          # set
          partitions = self._strategy.query_glob(MODEL_ID, "*")
          # logger.info(f"Fetching updates for {len(partitions)} models...")
          self._status_counter = Counter()
          for partition in partitions:
              last_status = partition.ctx.get(STATUS)
  
              # If we are done successfully, go on to the next one
              if last_status in (Status.COMPLETED,):
                  self._status_counter.update([Status(last_status)])
                  continue
  
              model_id = partition.ctx.get(MODEL_ID)
  
              # Hydrate a Model object from the remote API
              current_model = Model(self._project, model_id=model_id)
  
              # Update our strategy-wide counter of model states
              self._status_counter.update([current_model.status])
  
              # Did our model status change?
              if last_status != current_model.status:
                  logger.info(
                      f"Partition {partition.idx} status change from {last_status} to {current_model.status}"
                  )
  
              _update = {STATUS: current_model.status}
  
              if current_model.status == Status.COMPLETED:
                  report = current_model.peek_report()
  
                  if report is None:
                      with smart_open.open(
                          current_model.get_artifact_link("report_json")
                      ) as fin:
                          report = json.loads(fin.read())
  
                  sqs = report["synthetic_data_quality_score"]["score"]
                  label = "Moderate"
                  if sqs >= 80:
                      label = "Excellent"
                  elif sqs >= 60:
                      label = "Good"
  
                  if last_status != current_model.status:
                      logger.info(
                          f"Partition {partition.idx} completes with SQS: {label} ({sqs})"
                      )
  
                  _update.update({SQS: report})
  
              partition.update_ctx(_update)
              self._strategy.status_counter = dict(self._status_counter)
              # Aggressive, but save after every update
              self._strategy.status_counter = dict(self._status_counter)
              self._strategy.save()
  
          # If every partition is done, we may not have saved the strategy
          self._strategy.status_counter = dict(self._status_counter)
          self._strategy.save()
  
      def _update_handler_status(self):
          partitions = self._strategy.partitions
          self._handler_status_counter = Counter()
          for partition in partitions:
              model_id = partition.ctx.get(MODEL_ID)
              handler_id = partition.ctx.get(HANDLER, {}).get(HANDLER_ID)
              last_status = partition.ctx.get(HANDLER, {}).get(STATUS)
  
              # No need to refresh completed handlers from the remote API
              if last_status in (Status.COMPLETED,):
                  self._handler_status_counter.update([Status(last_status)])
                  continue
  
              if not handler_id:
                  continue
  
              # Hydrate a Model and Handler object from the remote API
              current_model = Model(self._project, model_id=model_id)
              current_handler = RecordHandler(
                  model=current_model, record_id=handler_id)
  
              self._handler_status_counter.update([current_handler.status])
  
              if last_status != current_handler.status:
                  logger.info(
                      f"Partition {partition.idx} record generation status change from {last_status} to {current_handler.status}"
                  )
  
              partition.ctx[HANDLER][STATUS] = current_handler.status
              self._strategy.save()
  
      @_needs_load
      def cancel_all(self):
          partitions = self._strategy.query_glob(MODEL_ID, "*")
          for partition in partitions:
              model_id = partition.ctx.get(MODEL_ID)
  
              # Hydrate a Model object from the remote API
              current_model = Model(self._project, model_id=model_id)
              logger.warning(f"Cancelling: {current_model.id}")
              current_model.cancel()
  
      def _refresh_max_job_capacity(self):
          self._max_jobs_active = get_me()["service_limits"]["max_jobs_active"]
  
      @property
      @_needs_load
      def has_capacity(self) -> bool:
          num_active = len(self._gather_statuses(ACTIVE_STATES))
          self._refresh_max_job_capacity()
          return num_active < self._max_jobs_active
  
      def _remove_unused_artifact(self) -> Optional[str]:
          project_artifacts = self._project.artifacts
          curr_artifacts = set()
  
          if len(project_artifacts) < self._max_artifacts:
              return "__none__"
  
          # First we try and remove an artifact from this strategy by
          # looking at our model states
          for p in self._strategy.partitions:
              status = p.ctx.get(STATUS)
              artifact_key = p.ctx.get(ARTIFACT)
  
              curr_artifacts.add(artifact_key)
  
              # We don't want to delete an artifact that is maybe
              # about to be used
              if status is None or status in ACTIVE_STATES:
                  continue
  
              if artifact_key:
                  logger.debug(
                      f"Attempting to remove artifact: {p.ctx.get(ARTIFACT)}")
                  self._project.delete_artifact(artifact_key)
                  p.update_ctx({ARTIFACT: None})
                  self._strategy.save()
                  return artifact_key
  
              # try and remove a seed artifact if one exists
              handler_status = p.ctx.get(HANDLER, {}).get(STATUS)
              handler_artifact_key = p.ctx.get(HANDLER, {}).get(ARTIFACT)
  
              if handler_status is None or handler_status in ACTIVE_STATES:
                  continue
  
              if handler_artifact_key:
                  logger.debug(
                      f"Attempting to remove handler artifact: {handler_artifact_key}"
                  )
                  self._project.delete_artifact(handler_artifact_key)
                  p.ctx[HANDLER][ARTIFACT] = None
                  return handler_artifact_key
  
          # If we couldn't remove an artifact from this current strategy,
          # we'll just remove some other random one
          try:
              logger.debug("Removing artifact not belonging to this Strategy...")
              for art in project_artifacts:
                  key = art.get("key")
                  if key in curr_artifacts:
                      continue
                  self._project.delete_artifact(key)
                  return key
          except Exception as err:
              logger.warning(f"Could not delete artifact: {str(err)}")
  
          return None
  
      def _partition_to_artifact(self, partition: Partition) -> Optional[ArtifactResult]:
          removed_artifact = self._remove_unused_artifact()
          if not removed_artifact:
              logger.debug(
                  "Could not make room for next data set, waiting for room...")
              # We couldn't make room so we don't try and upload the next artifact
              return None
  
          filename = f"{self.strategy_id}-{partition.idx}.csv"
          df_to_upload = partition.extract_df(self._df)
          res = self._df_to_artifact(df_to_upload, filename)
          partition.update_ctx({ARTIFACT: res.id})
          self._strategy.save()
          return res
  
      def _df_to_artifact(self, df: pd.DataFrame, filename: str) -> ArtifactResult:
          with tempfile.TemporaryDirectory() as tmp:
              target_file = str(Path(tmp) / filename)
              df.to_csv(target_file, index=False)
              artifact_id = self._project.upload_artifact(target_file)
              return ArtifactResult(id=artifact_id, record_count=len(df))
  
      @_needs_load
      def train_partition(
          self, partition: Partition, artifact: ArtifactResult
      ) -> Optional[str]:
          attempt = partition.ctx.get(ATTEMPT, 0) + 1
          model_config = deepcopy(self._model_config)
          data_source = None
  
          if "synthetics" in model_config["models"][0].keys():
  
              model_config["models"][0]["synthetics"]["generate"] = {
                  "num_records": artifact.record_count,
                  "max_invalid": None,
              }
  
              # If we're trying this model for a second+ time, we reduce the vocab size to
              # utilize the char encoder in order to give a better chance and success
              if attempt > 1:
                  model_config["models"][0]["synthetics"]["params"]["vocab_size"] = 0
  
              # If this partition is for the first-N headers and we have known seed headers, we have to
              # modify the configuration to account for the seed task.
              if partition.columns.seed_headers:
                  model_config["models"][0]["synthetics"]["task"] = {
                      "type": "seed",
                      "attrs": {"fields": partition.columns.seed_headers},
                  }
  
          elif "ctgan" in model_config["models"][0].keys():
              pass
  
          model = self._project.create_model_obj(
              model_config=model_config, data_source=artifact.id
          )
  
          model = _maybe_submit_job(model)
          if model is None:
              return None
  
          partition.ctx.update(
              {
                  STATUS: model.status,
                  ARTIFACT: artifact.id,
                  MODEL_ID: model.model_id,
                  ATTEMPT: attempt,
              }
          )
          self._strategy.save()
          logger.info(
              f"Started model: {model.print_obj['model_name']} " f"source: {artifact.id}"
          )
          return model.model_id
  
      @_needs_load
      def run_partition(
          self, partition: Partition, gen_payload: GenPayload
      ) -> Optional[str]:
          """
          Run a record handler for a model and return the job id.
  
          NOTE: This assumes the partition is successfully trained and has an
          available model.
          """
          handler_dict = partition.ctx.get(HANDLER)
          if handler_dict is None:
              partition.ctx[HANDLER] = {}
          attempt = partition.ctx.get(HANDLER).get(ATTEMPT, 0) + 1
          model_id = partition.ctx.get(MODEL_ID)
  
          # Hydrate our trained model so we can start the handler
          model_obj = Model(self._project, model_id=model_id)
  
          # Create and start our handler to generate data
          handler_obj = model_obj.create_record_handler_obj(
              data_source=gen_payload.seed_artifact_id,
              params={
                  "num_records": gen_payload.num_records,
                  "max_invalid": gen_payload.max_invalid,
              },
          )
          handler_obj = _maybe_submit_job(handler_obj)
          if handler_obj is None:
              return None
  
          _ctx_update = {
              ATTEMPT: attempt,
              ARTIFACT: gen_payload.seed_artifact_id,
              NUM_RECS: gen_payload.num_records,
              STATUS: handler_obj.status,
              HANDLER_ID: handler_obj.record_id,
          }
  
          partition.ctx[HANDLER].update(_ctx_update)
          self._strategy.save()
          logger.info(
              f"Generating {gen_payload.num_records} records from model: {model_obj.print_obj['model_name']}"
          )
          return handler_obj.record_id
  
      @_needs_load
      def train_next_partition(self) -> Optional[str]:
          start_job = False
          for partition in self._strategy.partitions:
              status = partition.ctx.get(STATUS)  # type: Status
  
              # If we've never done a job for this partition, we should start one
              if status is None:
                  logger.info(
                      f"Partition {partition.idx} is new, starting model creation"
                  )
                  start_job = True
  
              # If the job failed, should we try again?
              elif (
                  status in (Status.ERROR, Status.LOST)
                  and partition.ctx.get(ATTEMPT, 0) < self._error_retry_limit
              ):
                  logger.info(
                      f"Partition {partition.idx} status: {status.value}, re-attempting job"
                  )
                  start_job = True
  
              if start_job:
                  artifact = self._partition_to_artifact(partition)
                  if artifact.id is None:
                      return None
                  return self.train_partition(partition, artifact)
  
      @_needs_load
      def run_next_partition(self, gen_payload: GenPayload) -> Optional[str]:
          start_job = False
          for partition in self._strategy.partitions:
              status = partition.ctx.get(HANDLER, {}).get(STATUS)  # type: Status
              attempt_count = partition.ctx.get(HANDLER, {}).get(ATTEMPT, 0)
  
              if status is None:
                  logger.info(f"Generating data for partition {partition.idx}")
                  start_job = True
  
              elif (
                  status in (Status.ERROR, Status.LOST)
                  and attempt_count < self._error_retry_limit
              ):
                  logger.info(
                      f"Partition {partition.idx} has status {status.value}, re-attempting generation"
                  )
                  start_job = True
  
              if start_job:
                  use_seeds = False
                  # If this partition has seed fields and we were given seeds, we need to upload
                  # the artifact first.
                  if partition.columns.seed_headers and isinstance(
                      gen_payload.seed_df, pd.DataFrame
                  ):
                      # NOTE(jm): If we've tried N-1 attempts with seeds and the handler has continued
                      # fail then we should not use seeds to at least let the handler try to succeed.
                      # One example of this happening would be when a partition's model receives seeds
                      # where the values of the seeds were not in the training set (due to partitioning).
                      if attempt_count == self._error_retry_limit - 1:
                          logger.info(
                              f"WARNING: Disabling seeds for partition {partition.idx} due to previous failed generation attempts..."
                          )
                      else:
                          logger.info(
                              "Partition has seed fields, uploading seed artifact..."
                          )
                          use_seeds = True
                          removed_artifact = self._remove_unused_artifact()
                          if removed_artifact is None:
                              logger.info(
                                  "Could not start generation with seeds, an old artifact could not be removed"
                              )
                              return None
  
                          filename = f"{self.strategy_id}-seeds-{partition.idx}.csv"
                          artifact = self._df_to_artifact(
                              gen_payload.seed_df, filename)
  
                  new_payload = GenPayload(
                      num_records=gen_payload.num_records,
                      max_invalid=gen_payload.max_invalid,
                      seed_artifact_id=artifact.id if use_seeds else None,
                  )
  
                  return self.run_partition(partition, new_payload)
  
      @_needs_load
      def clear_partition_runs(self):
          """
          Partitions should only be trained until they are 'completed', however we can run
          a partition any number of times. Before we do that, we want to go through and
          """
          for partition in self._strategy.partitions:
              partition.ctx[HANDLER] = {}
  
      def _gather_statuses(self, statuses: List[Status]) -> List[Partition]:
          out = []
          for partition in self._strategy.partitions:
              status = partition.ctx.get(STATUS)
              if status is None:
                  continue
              if status in statuses:
                  out.append(partition)
          return out
  
      @_needs_load
      def is_done(self, *, handler: bool = False) -> bool:
          done = 0
          for p in self._strategy.partitions:
              if handler:
                  ctx_base = p.ctx.get(HANDLER, {})
              else:
                  ctx_base = p.ctx
  
              status = ctx_base.get(STATUS)
              attempt = ctx_base.get(ATTEMPT, 0)
              if status is None:
                  continue
  
              if status in (Status.COMPLETED, Status.CANCELLED):
                  done += 1
              elif (
                  status in (Status.ERROR, Status.LOST)
                  and attempt >= self._error_retry_limit
              ):
                  done += 1
  
          return done == len(self._strategy.partitions)
  
      @_needs_load
      def train_all_partitions(self):
          logger.info(f"Processing {len(self._strategy.partitions)} partitions")
          while True:
              self._update_job_status()
              if not self.has_capacity:
                  logger.debug("At active capacity, waiting for more...")
                  time.sleep(10)
                  continue
  
              model_id = self.train_next_partition()
              if model_id:
                  continue
  
              if self.is_done():
                  break
  
              time.sleep(10)
  
          logger.info(dict(self._status_counter))
  
      @_needs_load
      def _get_synthetic_data(self, job_type: str, artifact_type: str) -> pd.DataFrame:
          if job_type == "model":
              self._update_job_status()
              num_completed = self._status_counter.get(Status.COMPLETED, 0)
          elif job_type == "run":
              self._update_handler_status()
              num_completed = self._handler_status_counter.get(
                  Status.COMPLETED, 0)
          else:
              raise ValueError("invalid job_type")
  
          if num_completed != self._strategy.partition_count:
              raise RuntimeError(
                  "Not all partitions are completed, cannot fetch synthetic data from trained models"
              )
  
          # We will have at least one column-wise DF, this holds
          # one DF for each header cluster we have
          df_chunks = {
              i: pd.DataFrame() for i in range(0, self._strategy.header_cluster_count)
          }
  
          pool = ThreadPoolExecutor()
          futures = []
          for partition in self._strategy.partitions:
              # NOTE: When sending the remote DF payload we try to extract both
              # model and handler IDs and the thread workers can interpret which
              # ones they need to use.
              payload = RemoteDFPayload(
                  partition=partition.idx,
                  slot=partition.columns.idx,
                  job_type=job_type,
                  handler_uid=partition.ctx.get(HANDLER, {}).get(HANDLER_ID),
                  uid=partition.ctx.get(MODEL_ID),
                  project=self._project,
                  artifact_type=artifact_type,
              )
              futures.append(pool.submit(_remote_dataframe_fetcher, payload))
  
          wait(futures, return_when=ALL_COMPLETED)
  
          for future in futures:
              payload = future.result()  # type: RemoteDFPayload
  
              curr_df = df_chunks[payload.slot]
              df_chunks[payload.slot] = pd.concat([curr_df, payload.df]).reset_index(
                  drop=True
              )
  
          df = pd.concat(list(df_chunks.values()), axis=1)
          return df
  
      def _maybe_restore_df_headers(self, df) -> pd.DataFrame:
          if isinstance(self._df, pd.DataFrame):
              return df[self._df.columns]
  
          if self._strategy.original_headers:
              return df[self._strategy.original_headers]
  
          return df
  
      def get_training_synthetic_data(self) -> pd.DataFrame:
          df = self._get_synthetic_data("model", "data_preview")
          return self._maybe_restore_df_headers(df)
  
      def get_synthetic_data(self) -> pd.DataFrame:
          df = self._get_synthetic_data("run", "data")
          return self._maybe_restore_df_headers(df)
  
      @_needs_load
      def generate_data(
          self,
          *,
          seed_df: Optional[pd.DataFrame] = None,
          num_records: Optional[int] = None,
          max_invalid: Optional[int] = None,
          clear_cache: bool = False,
      ):
  
          if seed_df is None and not num_records:
              raise ValueError(
                  "must provide a seed_df or num_records to generate")
  
          if isinstance(seed_df, pd.DataFrame) and num_records:
              raise ValueError("must use one of seed_df or num_records only")
  
          # Refresh all of the trained models
          logger.info("Loading existing model information...")
          self._update_job_status()
  
          # We can't generate a dataset if any of the models are in a bad state, so we check that here
          completed_count = self._status_counter[Status.COMPLETED]
          if completed_count != self._strategy.partition_count:
              raise RuntimeError(
                  f"Cannot generate data, {self._strategy.partition_count - completed_count} partitions do not have a completed model."
              )
  
          # Clear out all previous record handler metadata
          if clear_cache:
              self.clear_partition_runs()
  
          # If we have seeds, then we use the number of seeds as the number of records
          # to generate from each model.
          found_seeds = False
          if isinstance(seed_df, pd.DataFrame):
              num_records = len(seed_df)
  
              # Loop through all of the partitions and make sure we have some that
              # take seed values, if we don't have any partitions set for seeds
              # and we recieved a seed DF, we should error.
              for partition in self._strategy.partitions:
                  if partition.columns.seed_headers:
                      found_seeds = True
                      break
  
              if not found_seeds:
                  raise RuntimeError(
                      "You cannot provide a seed_df since models were not conditioned with seed headers."
                  )
  
          # NOTE: This payload will be used to create a new payload object per
          # partition, so this will get passed in more as a template to the next
          # routine
          gen_payload = GenPayload(
              seed_df=seed_df, num_records=num_records, max_invalid=max_invalid
          )
  
          logger.info(
              f"Starting data generation from {self._strategy.partition_count} models"
          )
  
          update_remote = True
  
          while True:
              # Try and start as many jobs as we can, until we hit a rate
              # limit error, then we can check our capacity availabiliy and
              # wait until trying more.
              #
              # The main idea is that we want to avoid polling for re-hydration
              # so much since that can take so long.
              if update_remote:
                  logger.info("Refreshing generation statuses...")
                  self._update_handler_status()
                  update_remote = False
  
              try:
                  handler_id = self.run_next_partition(gen_payload)
                  if handler_id is not None:
                      # Go around and try again if this succeeded
                      continue
  
              # Catch a 4xx in the event we are at capacity, or something else goes wrong
              except Exception as err:
                  logger.warning(f"Error running model: {str(err)}")
  
              # Are we done done?
              if self.is_done(handler=True):
                  break
  
              # Wait until we have capacity, once we have capacity, we'll
              # set our flag to refresh the remote status update
              while True:
                  if not self.has_capacity:
                      logger.info("Waiting for more job capacity.")
                      time.sleep(10)
                  else:
                      update_remote = True
                      break
  
              time.sleep(10)
  
          logger.info(dict(self._handler_status_counter))
  