Esempio n. 1
0
class ModelSwapperInterface(object):
    """
  This is the interface class to connect the application layer to the Model
  Swapper.
  """

    #_INPUT_Q_OPTION_NAME = "input_queue"

    #_INPUT_Q_ENV_VAR = ModelSwapperConfig.getEnvVarOverrideName(
    #  configName=ModelSwapperConfig.CONFIG_NAME,
    #  section=_CONFIG_SECTION,
    #  option=_INPUT_Q_OPTION_NAME)
    #""" For testing: environment variable for overriding input queue name """

    _CONFIG_SECTION = "interface_bus"

    _RESULTS_Q_OPTION_NAME = "results_queue"

    # For testing: environment variable for overriding output queue name
    _RESULTS_Q_ENV_VAR = ModelSwapperConfig()._getEnvVarOverrideName(
        configName=ModelSwapperConfig.CONFIG_NAME,
        section=_CONFIG_SECTION,
        option=_RESULTS_Q_OPTION_NAME)

    _SCHEDULER_NOTIFICATION_Q_OPTION_NAME = "scheduler_notification_queue"

    _MODEL_INPUT_Q_PREFIX_OPTION_NAME = "model_input_queue_prefix"

    def __init__(self):
        """
    Initialize the ModelSwapperInterface. This uses a lazy loading of the input
    and output queues with no pre-meditation.
    """
        self._logger = _getLogger()

        config = ModelSwapperConfig()

        self._resultsQueueName = config.get(self._CONFIG_SECTION,
                                            self._RESULTS_Q_OPTION_NAME)

        # The name of a model's input message queue is the concatenation of this
        # prefix and the modelID
        self._modelInputQueueNamePrefix = config.get(
            self._CONFIG_SECTION, self._MODEL_INPUT_Q_PREFIX_OPTION_NAME)

        self._schedulerNotificationQueueName = config.get(
            self._CONFIG_SECTION, self._SCHEDULER_NOTIFICATION_Q_OPTION_NAME)

        # Message bus connector
        self._bus = MessageBusConnector()

        # Outstanding request and/or response consumer instances
        self._consumers = []

    def __enter__(self):
        """ Context Manager protocol method. Allows a ModelSwapperInterface instance
    to be used in a "with" statement for automatic clean-up

    Parameters:
    ------------------------------------------------------------------------
    retval:     self.
    """
        return self

    def __exit__(self, _excType, _excVal, _excTb):
        """ Context Manager protocol method. Allows a ModelSwapperInterface instance
    to be used in a "with" statement for automatic cleanup

    :returns: False so as not to suppress the exception, if any
    """
        self.close()
        return False

    def close(self):
        """
    Gracefully close ModelSwapperInterface instance (e.g., tear down
    connections). If this is not called, the underlying connections will
    eventually timeout, but it is good practice to close explicitly.
    """
        if self._consumers:
            self._logger.error(
                "While closing %s, discovered %s unclosed consumers; will "
                "attempt to close them now", self.__class__.__name__,
                len(self._consumers))

            for consumer in tuple(self._consumers):
                consumer.close()

            assert not self._consumers

        try:
            self._bus.close()
        finally:
            self._bus = None

    def _onConsumerClosed(self, consumer):
        """ Called by consumer instance's close() method to remove the consumer from
    our outstanding consumers list
    """
        self._consumers.remove(consumer)

    def _getModelInputQName(self, modelID):
        return self._modelInputQueueNamePrefix + modelID

    def _getModelIDFromInputQName(self, mqName):
        assert mqName.startswith(self._modelInputQueueNamePrefix), (
            "mq=%s doesn't start with %s") % (mqName,
                                              self._modelInputQueueNamePrefix)

        return mqName[len(self._modelInputQueueNamePrefix):]

    def defineModel(self, modelID, args, commandID):
        """ Initialize model's input message queue and send the "defineModel"
    command. The ModelCommandResult will be delivered asynchronously, along with
    the corresponding commandID and no args, to the process that is consuming
    ModelSwapper results.

    :param modelID: a hex string that uniquely identifies the target model.
    :param args: dict with the following properties:
      "modelConfig": model config dict suitable for passing to OPF
        ModelFactory.create()
      "inferenceArgs": Model inference arguments suitable for passing to
        model.enableInference()
      "inputRecordSchema": a sequence  of nupic.data.fieldmeta.FieldMetaInfo
        instances with field names/types/special as expected by the model and in
        the same order as they will appear in input records. This is needed in
        order to avoid the overhead of passing fields names with each and every
        input record, while permitting the necessary dictionaries to be
        constructed by ModelRunner for input to the OPF Model.
    :param commandID: a numeric or string id to associate with the command and
      result.
    """
        # TODO: validate input args dict against schema

        mqName = self._getModelInputQName(modelID)

        self._bus.createMessageQueue(mqName, durable=True)

        self.submitRequests(modelID,
                            (ModelCommand(commandID, "defineModel", args), ))

    def cloneModel(self, modelID, newModelID, commandID):
        """ Initiate cloning of an existing model. Initialize the new
    model's input message queue and send the "cloneModel" command to the source
    model. The ModelCommandResult will be delivered asynchronously, along with
    the corresponding commandID and no args, to the process that is consuming
    ModelSwapper results.

    :param modelID: a hex string that uniquely identifies the existing model.
    :param newModelID: a hex string that uniquely identifies the new model.
    :param commandID: a numeric or string id to associate with the command and
      result.

    :raises: ModelNotFound if the source model's input endpoint doesn't exist
    """
        # Create the model input message queue for the new model
        self._bus.createMessageQueue(self._getModelInputQName(newModelID),
                                     durable=True)

        self.submitRequests(modelID, (ModelCommand(
            commandID, "cloneModel", args={"modelID": newModelID}), ))

    def deleteModel(self, modelID, commandID):
        """ Submit a request to delete a model. The ModelCommandResult will be
    delivered asynchronously, along with the corresponding commandID and no
    args, to the process that is consuming ModelSwapper results.

    This method is idempotent.

    :param modelID: a hex string that uniquely identifies the target model.
    :param commandID: a numeric or string id to associate with the command and
                      result.
    """
        # First, purge unread input messages for this model, if any, to avoid
        # unnecessary processing before the model is deleted
        mq = self._getModelInputQName(modelID)
        self._logger.info(
            "deleteModel: purging mq=%s before submitting "
            "deleteModel command for model=%s", mq, modelID)
        try:
            self._bus.purge(mq)
        except message_bus_connector.MessageQueueNotFound:
            # deleteModel is an idempotent operation: assume this exception is
            # due to repeated attempt
            pass
        else:
            try:
                self.submitRequests(modelID,
                                    (ModelCommand(commandID, "deleteModel"), ))
            except ModelNotFound:
                # deleteModel is an idempotent operation: assume this exception is
                # due to repeated attempt
                pass

    def cleanUpAfterModelDeletion(self, modelID):
        """ For use by Engine's ModelRunner after it deletes a model: clean up
    resources that ModelSwapperInterface created to support the model, such
    as deleting the model's input message queue
    """
        self._bus.deleteMessageQueue(self._getModelInputQName(modelID))

    def modelInputPending(self, modelID):
        """ Check if input requests are pending for a model

    :param modelID: a string that uniquely identifies the target model.

    :returns: True if the model's input queue exists and is non-empty;
              False if the model's input queue is non-empty or doesn't exist
    """
        try:
            return not self._bus.isEmpty(self._getModelInputQName(modelID))
        except message_bus_connector.MessageQueueNotFound:
            return False

    def getModelsWithInputPending(self):
        """ Get model IDs of all models with pending input (non-empty input queues)

    :returns: (possibly empty) sequence of model IDs whose input streams are
      non-empty
    """

        # NOTE: queues may be deleted as we're running through the list, so we need
        # to play it safe
        def safeIsInputPending(mq):
            try:
                return not self._bus.isEmpty(mq)
            except message_bus_connector.MessageQueueNotFound:
                return False

        prefix = self._modelInputQueueNamePrefix
        return tuple(
            self._getModelIDFromInputQName(mq)
            for mq in self._bus.getAllMessageQueues()
            if mq.startswith(prefix) and safeIsInputPending(mq))

    def submitRequests(self, modelID, requests):
        """
    Submit a batch of requests for processing by a model with the given modelID.

    NOTE: it's an error to submit requests for a model after calling
    deleteModel()

    Keyword arguments:
    :param modelID: a string that uniquely identifies the target model.

    :param requests: a sequence of ModelCommand and/or ModelInputRow instances.
      NOTE: To create or delete a model, call the createModel or deleteModel
      method instead of submitting the "defineModel" or "deleteModel" commands.
      Together, the sequence of requests constitutes a request "batch".

    :returns: UUID of the submitted batch (intended for test code only)

    :raises: ModelNotFound if model's input endpoint doesn't exist

    Requests for a specific model will be processed in the submitted order.
    The results will be delivered asynchronously, along with the corresponding
    requestIDs, to the process that is consuming ModelSwapper results.

    NOTE: This assumes retry logic will be handled by the underlying MQ
    implementation.
    """
        batchID = uuid.uuid1().hex
        msg = RequestMessagePackager.marshal(
            batchID=batchID, batchState=BatchPackager.marshal(batch=requests))

        mqName = self._getModelInputQName(modelID)
        try:
            self._bus.publish(mqName, msg, persistent=True)
        except message_bus_connector.MessageQueueNotFound as e:
            self._logger.warn(
                "App layer attempted to submit numRequests=%s to model=%s, but its "
                "input queue doesn't exist. Likely a race condition with model "
                "deletion path.", len(requests), modelID)
            raise ModelNotFound(repr(e))
        except:
            self._logger.exception(
                "Failed to publish request batch=%s for model=%s via mq=%s; "
                "msgLen=%s; msgPrefix=%r. NOTE: it's an error to submit requests to a "
                "model after deleting it.", batchID, modelID, mqName, len(msg),
                msg[:32])
            raise

        # Send a notification to Model Scheduler so it will schedule the model
        # for processing input
        try:
            self._bus.publish(self._schedulerNotificationQueueName,
                              json.dumps(modelID),
                              persistent=False)
        except message_bus_connector.MessageQueueNotFound:
            # If it's not fully up yet, its notification queue might not have been
            # created, which is ok
            self._logger.warn(
                "Couldn't send model data notification to Model Scheduler: mq=%s not "
                "found. Model Scheduler service not started or initialized the mq yet?",
                self._schedulerNotificationQueueName)
        return batchID

    def consumeRequests(self, modelID, blocking=True):
        """ Create an instance of the _MessageConsumer iterable for reading model
    requests, a batch at a time. The iterable yields _ConsumedRequestBatch
    instances.

    NOTE: This API is intended for Engine Model Runners.

    :param modelID: a string that uniquely identifies the target model.
    :param blocking: if True, the iterable will block until another batch becomes
      available; if False, the iterable will terminate iteration when no more
      batches are available in the queue. [defaults to True]

    :returns: an instance of model_swapper_interface._MessageConsumer iterable;
      IMPORTANT: the caller is responsible for closing it before closing this
      ModelSwapperInterface instance (hint: use the returned _MessageConsumer
      instance as Context Manager)

    :raises: ModelNotFound if model's input endpoint doesn't exist
            TODO: need tests for consumeRequests with ModelNotFound

    Example:
      with ModelSwapperInterface() as swapper:
        with swapper.consumeRequests(modelID) as consumer:
          for batch in consumer:
            processRequests(batchID=batch.batchID, requests=batch.objects)
            batch.ack()
    """
        mq = self._getModelInputQName(modelID)

        def onQueueNotFound():
            msg = (
                "Attempt to consume requests from model=%s is impossible because "
                "its input queue doesn't exist. Likely a race condition with "
                "model deletion path.") % (modelID, )
            self._logger.warn(msg)
            raise ModelNotFound(msg)

        consumer = _MessageConsumer(mqName=mq,
                                    blocking=blocking,
                                    decode=_ConsumedRequestBatch.decodeMessage,
                                    swapper=self,
                                    bus=self._bus,
                                    onQueueNotFound=onQueueNotFound)

        self._consumers.append(consumer)

        return consumer

    def _initResultsMessageQueue(self):
        self._bus.createMessageQueue(self._resultsQueueName, durable=True)

    def submitResults(self, modelID, results):
        """
    Submit a batch of results (used by ModelSwapper layer)

    Keyword arguments:
    :param modelID: a string that uniquely identifies the target model.

    :param results: a sequence of ModelCommandResult and/or ModelInferenceResult
      instances

    NOTE: This assumes retry logic will be handled by the underlying MQ
    implementation.
    """
        msg = ResultMessagePackager.marshal(
            modelID=modelID, batchState=BatchPackager.marshal(batch=results))
        try:
            try:
                self._bus.publish(self._resultsQueueName, msg, persistent=True)
            except message_bus_connector.MessageQueueNotFound:
                self._logger.info(
                    "submitResults: results mq=%s didn't exist; "
                    "declaring now and re-publishing message",
                    self._resultsQueueName)
                self._initResultsMessageQueue()
                self._bus.publish(self._resultsQueueName, msg, persistent=True)
        except:
            self._logger.exception(
                "submitResults: Failed to publish results from model=%s via mq=%s; "
                "msgLen=%s; msgPrefix=%r", modelID, self._resultsQueueName,
                len(msg), msg[:32])
            raise

    def consumeResults(self):
        """ Create an instance of the _MessageConsumer iterable for reading model
    results, a batch at a time. The iterable yields _ConsumedResultBatch
    instances.

    :returns: an instance of model_swapper_interface._MessageConsumer iterable;
      IMPORTANT: the caller is responsible for closing it before closing this
      ModelSwapperInterface instance (hint: use the returned _MessageConsumer
      instance as Context Manager)

    Example:
      with ModelSwapperInterface() as swapper:
        with swapper.consumeResults() as consumer:
          for batch in consumer:
            processResults(modelID=batch.modelID, results=batch.objects)
            batch.ack()
    """
        consumer = _MessageConsumer(
            mqName=self._resultsQueueName,
            blocking=True,
            decode=_ConsumedResultBatch.decodeMessage,
            swapper=self,
            bus=self._bus,
            onQueueNotFound=self._initResultsMessageQueue)

        self._consumers.append(consumer)

        return consumer

    def initSchedulerNotification(self):
        """ Initialize Model Scheduler's notification message queue; for use by
    Model Scheduler.
    """
        self._bus.createMessageQueue(self._schedulerNotificationQueueName,
                                     durable=False)

    def consumeModelSchedulerNotifications(self):
        """ Create an instance of the _MessageConsumer iterable for reading model
    scheduler notifications. The iterable yields _ConsumedNotification
    instances.

    :returns: an instance of model_swapper_interface._MessageConsumer iterable;
      IMPORTANT: the caller is responsible for closing it before closing this
      ModelSwapperInterface instance (hint: use the returned _MessageConsumer
      instance as Context Manager)

    Example:
      with ModelSwapperInterface() as swapper:
        with swapper.consumeModelSchedulerNotifications() as consumer:
          for notification in consumer:
            processNotification(notification.value)
            notification.ack()
    """
        consumer = _MessageConsumer(
            mqName=self._schedulerNotificationQueueName,
            blocking=True,
            decode=_ConsumedNotification.decodeMessage,
            swapper=self,
            bus=self._bus)

        self._consumers.append(consumer)

        return consumer
class ModelSwapperInterface(object):
  """
  This is the interface class to connect the application layer to the Model
  Swapper.
  """

  #_INPUT_Q_OPTION_NAME = "input_queue"

  #_INPUT_Q_ENV_VAR = ModelSwapperConfig.getEnvVarOverrideName(
  #  configName=ModelSwapperConfig.CONFIG_NAME,
  #  section=_CONFIG_SECTION,
  #  option=_INPUT_Q_OPTION_NAME)
  #""" For testing: environment variable for overriding input queue name """

  _CONFIG_SECTION = "interface_bus"

  _RESULTS_Q_OPTION_NAME = "results_queue"

  # For testing: environment variable for overriding output queue name
  _RESULTS_Q_ENV_VAR = ModelSwapperConfig()._getEnvVarOverrideName(
    configName=ModelSwapperConfig.CONFIG_NAME,
    section=_CONFIG_SECTION,
    option=_RESULTS_Q_OPTION_NAME)

  _SCHEDULER_NOTIFICATION_Q_OPTION_NAME = "scheduler_notification_queue"

  _MODEL_INPUT_Q_PREFIX_OPTION_NAME = "model_input_queue_prefix"


  def __init__(self):
    """
    Initialize the ModelSwapperInterface. This uses a lazy loading of the input
    and output queues with no pre-meditation.
    """
    self._logger = _getLogger()

    config = ModelSwapperConfig()

    self._resultsQueueName = config.get(
      self._CONFIG_SECTION, self._RESULTS_Q_OPTION_NAME)

    # The name of a model's input message queue is the concatenation of this
    # prefix and the modelID
    self._modelInputQueueNamePrefix = config.get(
      self._CONFIG_SECTION, self._MODEL_INPUT_Q_PREFIX_OPTION_NAME)

    self._schedulerNotificationQueueName = config.get(
      self._CONFIG_SECTION, self._SCHEDULER_NOTIFICATION_Q_OPTION_NAME)

    # Message bus connector
    self._bus = MessageBusConnector()

    # Outstanding request and/or response consumer instances
    self._consumers = []


  def __enter__(self):
    """ Context Manager protocol method. Allows a ModelSwapperInterface instance
    to be used in a "with" statement for automatic clean-up

    Parameters:
    ------------------------------------------------------------------------
    retval:     self.
    """
    return self


  def __exit__(self, _excType, _excVal, _excTb):
    """ Context Manager protocol method. Allows a ModelSwapperInterface instance
    to be used in a "with" statement for automatic cleanup

    :returns: False so as not to suppress the exception, if any
    """
    self.close()
    return False


  def close(self):
    """
    Gracefully close ModelSwapperInterface instance (e.g., tear down
    connections). If this is not called, the underlying connections will
    eventually timeout, but it is good practice to close explicitly.
    """
    if self._consumers:
      self._logger.error(
        "While closing %s, discovered %s unclosed consumers; will "
        "attempt to close them now", self.__class__.__name__,
        len(self._consumers))

      for consumer in tuple(self._consumers):
        consumer.close()

      assert not self._consumers

    try:
      self._bus.close()
    finally:
      self._bus = None


  def _onConsumerClosed(self, consumer):
    """ Called by consumer instance's close() method to remove the consumer from
    our outstanding consumers list
    """
    self._consumers.remove(consumer)


  def _getModelInputQName(self, modelID):
    return self._modelInputQueueNamePrefix + modelID


  def _getModelIDFromInputQName(self, mqName):
    assert mqName.startswith(self._modelInputQueueNamePrefix), (
      "mq=%s doesn't start with %s") % (mqName, self._modelInputQueueNamePrefix)

    return mqName[len(self._modelInputQueueNamePrefix):]


  def defineModel(self, modelID, args, commandID):
    """ Initialize model's input message queue and send the "defineModel"
    command. The ModelCommandResult will be delivered asynchronously, along with
    the corresponding commandID and no args, to the process that is consuming
    ModelSwapper results.

    :param modelID: a hex string that uniquely identifies the target model.
    :param args: dict with the following properties:
      "modelConfig": model config dict suitable for passing to OPF
        ModelFactory.create()
      "inferenceArgs": Model inference arguments suitable for passing to
        model.enableInference()
      "inputRecordSchema": a sequence  of nupic.data.fieldmeta.FieldMetaInfo
        instances with field names/types/special as expected by the model and in
        the same order as they will appear in input records. This is needed in
        order to avoid the overhead of passing fields names with each and every
        input record, while permitting the necessary dictionaries to be
        constructed by ModelRunner for input to the OPF Model.
    :param commandID: a numeric or string id to associate with the command and
      result.
    """
    # TODO: validate input args dict against schema

    mqName = self._getModelInputQName(modelID)

    self._bus.createMessageQueue(mqName, durable=True)

    self.submitRequests(modelID,
                        (ModelCommand(commandID, "defineModel", args),))


  def cloneModel(self, modelID, newModelID, commandID):
    """ Initiate cloning of an existing model. Initialize the new
    model's input message queue and send the "cloneModel" command to the source
    model. The ModelCommandResult will be delivered asynchronously, along with
    the corresponding commandID and no args, to the process that is consuming
    ModelSwapper results.

    :param modelID: a hex string that uniquely identifies the existing model.
    :param newModelID: a hex string that uniquely identifies the new model.
    :param commandID: a numeric or string id to associate with the command and
      result.

    :raises: ModelNotFound if the source model's input endpoint doesn't exist
    """
    # Create the model input message queue for the new model
    self._bus.createMessageQueue(self._getModelInputQName(newModelID),
                                 durable=True)

    self.submitRequests(
      modelID,
      (ModelCommand(commandID, "cloneModel", args={"modelID": newModelID}),))


  def deleteModel(self, modelID, commandID):
    """ Submit a request to delete a model. The ModelCommandResult will be
    delivered asynchronously, along with the corresponding commandID and no
    args, to the process that is consuming ModelSwapper results.

    This method is idempotent.

    :param modelID: a hex string that uniquely identifies the target model.
    :param commandID: a numeric or string id to associate with the command and
                      result.
    """
    # First, purge unread input messages for this model, if any, to avoid
    # unnecessary processing before the model is deleted
    mq = self._getModelInputQName(modelID)
    self._logger.info("deleteModel: purging mq=%s before submitting "
                      "deleteModel command for model=%s", mq, modelID)
    try:
      self._bus.purge(mq)
    except message_bus_connector.MessageQueueNotFound:
      # deleteModel is an idempotent operation: assume this exception is
      # due to repeated attempt
      pass
    else:
      try:
        self.submitRequests(modelID, (ModelCommand(commandID, "deleteModel"),))
      except ModelNotFound:
        # deleteModel is an idempotent operation: assume this exception is
        # due to repeated attempt
        pass


  def cleanUpAfterModelDeletion(self, modelID):
    """ For use by Engine's ModelRunner after it deletes a model: clean up
    resources that ModelSwapperInterface created to support the model, such
    as deleting the model's input message queue
    """
    self._bus.deleteMessageQueue(self._getModelInputQName(modelID))


  def modelInputPending(self, modelID):
    """ Check if input requests are pending for a model

    :param modelID: a string that uniquely identifies the target model.

    :returns: True if the model's input queue exists and is non-empty;
              False if the model's input queue is non-empty or doesn't exist
    """
    try:
      return not self._bus.isEmpty(self._getModelInputQName(modelID))
    except message_bus_connector.MessageQueueNotFound:
      return False


  def getModelsWithInputPending(self):
    """ Get model IDs of all models with pending input (non-empty input queues)

    :returns: (possibly empty) sequence of model IDs whose input streams are
      non-empty
    """
    # NOTE: queues may be deleted as we're running through the list, so we need
    # to play it safe
    def safeIsInputPending(mq):
      try:
        return not self._bus.isEmpty(mq)
      except message_bus_connector.MessageQueueNotFound:
        return False

    prefix = self._modelInputQueueNamePrefix
    return tuple(
      self._getModelIDFromInputQName(mq)
      for mq in self._bus.getAllMessageQueues()
      if mq.startswith(prefix) and safeIsInputPending(mq))


  def submitRequests(self, modelID, requests):
    """
    Submit a batch of requests for processing by a model with the given modelID.

    NOTE: it's an error to submit requests for a model after calling
    deleteModel()

    Keyword arguments:
    :param modelID: a string that uniquely identifies the target model.

    :param requests: a sequence of ModelCommand and/or ModelInputRow instances.
      NOTE: To create or delete a model, call the createModel or deleteModel
      method instead of submitting the "defineModel" or "deleteModel" commands.
      Together, the sequence of requests constitutes a request "batch".

    :returns: UUID of the submitted batch (intended for test code only)

    :raises: ModelNotFound if model's input endpoint doesn't exist

    Requests for a specific model will be processed in the submitted order.
    The results will be delivered asynchronously, along with the corresponding
    requestIDs, to the process that is consuming ModelSwapper results.

    NOTE: This assumes retry logic will be handled by the underlying MQ
    implementation.
    """
    batchID = uuid.uuid1().hex
    msg = RequestMessagePackager.marshal(
      batchID=batchID,
      batchState=BatchPackager.marshal(batch=requests))

    mqName = self._getModelInputQName(modelID)
    try:
      self._bus.publish(mqName, msg, persistent=True)
    except message_bus_connector.MessageQueueNotFound as e:
      self._logger.warn(
        "App layer attempted to submit numRequests=%s to model=%s, but its "
        "input queue doesn't exist. Likely a race condition with model "
        "deletion path.", len(requests), modelID)
      raise ModelNotFound(repr(e))
    except:
      self._logger.exception(
        "Failed to publish request batch=%s for model=%s via mq=%s; "
        "msgLen=%s; msgPrefix=%r. NOTE: it's an error to submit requests to a "
        "model after deleting it.", batchID, modelID, mqName, len(msg),
        msg[:32])
      raise

    # Send a notification to Model Scheduler so it will schedule the model
    # for processing input
    try:
      self._bus.publish(self._schedulerNotificationQueueName,
                        json.dumps(modelID), persistent=False)
    except message_bus_connector.MessageQueueNotFound:
      # If it's not fully up yet, its notification queue might not have been
      # created, which is ok
      self._logger.warn(
        "Couldn't send model data notification to Model Scheduler: mq=%s not "
        "found. Model Scheduler service not started or initialized the mq yet?",
        self._schedulerNotificationQueueName)
    return batchID


  def consumeRequests(self, modelID, blocking=True):
    """ Create an instance of the _MessageConsumer iterable for reading model
    requests, a batch at a time. The iterable yields _ConsumedRequestBatch
    instances.

    NOTE: This API is intended for Engine Model Runners.

    :param modelID: a string that uniquely identifies the target model.
    :param blocking: if True, the iterable will block until another batch
      becomes available; if False, the iterable will terminate iteration when
      no more batches are available in the queue. [defaults to True]

    :returns: an instance of model_swapper_interface._MessageConsumer iterable;
      IMPORTANT: the caller is responsible for closing it before closing this
      ModelSwapperInterface instance (hint: use the returned _MessageConsumer
      instance as Context Manager)

    :raises: ModelNotFound if model's input endpoint doesn't exist
            TODO: need tests for consumeRequests with ModelNotFound

    Example:
      with ModelSwapperInterface() as swapper:
        with swapper.consumeRequests(modelID) as consumer:
          for batch in consumer:
            processRequests(batchID=batch.batchID, requests=batch.objects)
            batch.ack()
    """
    mq = self._getModelInputQName(modelID)

    def onQueueNotFound():
      msg = ("Attempt to consume requests from model=%s is impossible because "
             "its input queue doesn't exist. Likely a race condition with "
             "model deletion path.") % (modelID,)
      self._logger.warn(msg)
      raise ModelNotFound(msg)

    consumer = _MessageConsumer(mqName=mq,
                                blocking=blocking,
                                decode=_ConsumedRequestBatch.decodeMessage,
                                swapper=self,
                                bus=self._bus,
                                onQueueNotFound=onQueueNotFound)

    self._consumers.append(consumer)

    return consumer


  def _initResultsMessageQueue(self):
    self._bus.createMessageQueue(self._resultsQueueName, durable=True)


  def submitResults(self, modelID, results):
    """
    Submit a batch of results (used by ModelSwapper layer)

    Keyword arguments:
    :param modelID: a string that uniquely identifies the target model.

    :param results: a sequence of ModelCommandResult and/or ModelInferenceResult
      instances

    NOTE: This assumes retry logic will be handled by the underlying MQ
    implementation.
    """
    msg = ResultMessagePackager.marshal(
      modelID=modelID,
      batchState=BatchPackager.marshal(batch=results))
    try:
      try:
        self._bus.publish(self._resultsQueueName, msg, persistent=True)
      except message_bus_connector.MessageQueueNotFound:
        self._logger.info("submitResults: results mq=%s didn't exist; "
                          "declaring now and re-publishing message",
                          self._resultsQueueName)
        self._initResultsMessageQueue()
        self._bus.publish(self._resultsQueueName, msg, persistent=True)
    except:
      self._logger.exception(
        "submitResults: Failed to publish results from model=%s via mq=%s; "
        "msgLen=%s; msgPrefix=%r", modelID, self._resultsQueueName, len(msg),
        msg[:32])
      raise


  def consumeResults(self):
    """ Create an instance of the _MessageConsumer iterable for reading model
    results, a batch at a time. The iterable yields _ConsumedResultBatch
    instances.

    :returns: an instance of model_swapper_interface._MessageConsumer iterable;
      IMPORTANT: the caller is responsible for closing it before closing this
      ModelSwapperInterface instance (hint: use the returned _MessageConsumer
      instance as Context Manager)

    Example:
      with ModelSwapperInterface() as swapper:
        with swapper.consumeResults() as consumer:
          for batch in consumer:
            processResults(modelID=batch.modelID, results=batch.objects)
            batch.ack()
    """
    consumer = _MessageConsumer(mqName=self._resultsQueueName,
                                blocking=True,
                                decode=_ConsumedResultBatch.decodeMessage,
                                swapper=self,
                                bus=self._bus,
                                onQueueNotFound=self._initResultsMessageQueue)

    self._consumers.append(consumer)

    return consumer


  def initSchedulerNotification(self):
    """ Initialize Model Scheduler's notification message queue; for use by
    Model Scheduler.
    """
    self._bus.createMessageQueue(self._schedulerNotificationQueueName,
                                 durable=False)


  def consumeModelSchedulerNotifications(self):
    """ Create an instance of the _MessageConsumer iterable for reading model
    scheduler notifications. The iterable yields _ConsumedNotification
    instances.

    :returns: an instance of model_swapper_interface._MessageConsumer iterable;
      IMPORTANT: the caller is responsible for closing it before closing this
      ModelSwapperInterface instance (hint: use the returned _MessageConsumer
      instance as Context Manager)

    Example:
      with ModelSwapperInterface() as swapper:
        with swapper.consumeModelSchedulerNotifications() as consumer:
          for notification in consumer:
            processNotification(notification.value)
            notification.ack()
    """
    consumer = _MessageConsumer(mqName=self._schedulerNotificationQueueName,
                                blocking=True,
                                decode=_ConsumedNotification.decodeMessage,
                                swapper=self,
                                bus=self._bus)

    self._consumers.append(consumer)

    return consumer
  def testMessageBusConnectorCleanupOfConsumerGenerators(self):
    # Verify that MessageBusConnector closes unclosed consumers
    numMessagesToPublish = 1

    mqName = self._getUniqueMessageQueueName()

    with amqp_test_utils.managedQueueDeleter(mqName):
      with MessageBusConnector() as bus:
        # Create the queue
        bus.createMessageQueue(mqName=mqName, durable=True)

        # Now add messages
        expectedContent = []

        for i in xrange(numMessagesToPublish):
          expectedContent.append(str(i))
          bus.publish(mqName, expectedContent[-1], persistent=True)

        # Verify that correct number of messages were published
        self.assertEqual(_getQueueMessageCount(mqName), numMessagesToPublish)

      bus = MessageBusConnector()

      # Now, create a consumer iterable and start it by consuming the expected
      # messages
      # NOTE: we use a thread to avoid deadlocking the test runner in case
      #  something is wrong with the iterable
      def runConsumerThread(bus, mqName, numMessages, resultQ):
        try:
          # NOTE: in this test, we intentionally don't close the consumer
          consumer = bus.consume(mqName=mqName)
          resultQ.put(consumer)
          it = iter(consumer)
          for _i in xrange(numMessages):
            msg = next(it)
            msg.ack()
            resultQ.put(msg.body)
        except:
          resultQ.put(dict(exception=sys.exc_info()[1]))
          raise

      resultQ = Queue.Queue()
      consumerThread = threading.Thread(
        target=runConsumerThread,
        args=(bus, mqName, numMessagesToPublish, resultQ))
      consumerThread.setDaemon(True)
      consumerThread.start()

      # Wait for thread to stop
      consumerThread.join(timeout=30)
      self.assertFalse(consumerThread.isAlive())

      # Reap the consumer
      consumer1 = resultQ.get_nowait()
      self.assertIsInstance(consumer1, message_bus_connector._QueueConsumer)

      # Verify content
      actualContent = []
      while True:
        try:
          actualContent.append(resultQ.get_nowait())
        except Queue.Empty:
          break

      self.assertEqual(actualContent, expectedContent)

      # Verify that this consumer isn't closed
      self.assertEqual(len(bus._consumers), 1)
      self.assertIn(consumer1, bus._consumers)
      self.assertIsNotNone(consumer1._channelMgr)

      # Create another consumer, but don't start it (by telling it to consume 0
      # messages; we want to test cleanup of both started and unstarted
      # consumers
      resultQ = Queue.Queue()
      consumerThread = threading.Thread(
        target=runConsumerThread,
        args=(bus, mqName, 0, resultQ))
      consumerThread.setDaemon(True)
      consumerThread.start()

      # Wait for thread to stop
      consumerThread.join(timeout=30)
      self.assertFalse(consumerThread.isAlive())

      # Reap the consumer
      consumer2 = resultQ.get_nowait()
      self.assertIsInstance(consumer2, message_bus_connector._QueueConsumer)

      try:
        item = resultQ.get_nowait()
      except Queue.Empty:
        pass
      else:
        self.fail("Unexpected item in resultQ: %r" % (item,))

      # Verify that this consumer isn't closed
      self.assertEqual(len(bus._consumers), 2)
      self.assertIn(consumer2, bus._consumers)
      self.assertIsNotNone(consumer2._channelMgr)

      # Verify that this consumer hasn't started
      self.assertIsNone(consumer2._channelMgr._client)

      # Verify that the first consumer is still there, too
      self.assertIn(consumer1, bus._consumers)


      # Finanly, close the MessageBusConnector and verify that both consumers
      # got closed, too
      bus.close()

      self.assertFalse(bus._consumers)

      self.assertIsNone(consumer1._channelMgr)
      self.assertIsNone(consumer1._bus)
      self.assertIsNone(consumer1._channelMgr)

      self.assertIsNone(consumer2._channelMgr)
      self.assertIsNone(consumer2._bus)
      self.assertIsNone(consumer2._channelMgr)