def pull_logs()

in src/sagemaker/serve/utils/logging_agent.py [0:0]

• 50 lines of code
• 18 McCabe index (conditional complexity)

def pull_logs(generator, stop, until, final_pull):
    """Placeholder docstring"""
    now = datetime.now()
    if until < now:
        return

    logs = queue.Queue(maxsize=10)
    log_puller = Thread(
        target=_get_logs,
        args=(
            generator,
            logs,
            until,
        ),
    )
    log_puller.start()

    while True:
        try:
            top = logs.get(block=True, timeout=(until - now).total_seconds())
            logger.debug(top)

            if "[INFO ]" in top and "OutOfMemoryError" in top:
                raise LocalModelOutOfMemoryException(top)
            if "CUDA out of memory. Tried to allocate" in top:
                raise LocalModelOutOfMemoryException(top)
            if "ai.djl.engine.EngineException: OOM" in top:
                raise LocalModelOutOfMemoryException(top)
            if "4xx.Count:1" in top or "5xx.Count:1" in top:
                raise LocalModelInvocationException(top)
            if "[ERROR]" in top or "Failed register workflow" in top:
                raise LocalModelLoadException(top)
            if "Address already in use" in top:
                raise LocalModelLoadException(top)
            if "not compatible with sharding" in top:
                raise SkipTuningComboException(top)
        except (
            LocalModelLoadException,
            LocalModelOutOfMemoryException,
            LocalModelInvocationException,
        ) as e:
            stop()
            log_puller.join(timeout=5)
            raise e
        except queue.Empty:
            now = datetime.now()
            if until < now:
                if not final_pull:
                    return

                stop()
                # allow logging thread some time to cleanup
                log_puller.join(timeout=5)
                if log_puller.is_alive():
                    raise Exception("Logging thread not terminating")
                return