Display the logs for a given training job, optionally tailing them until the
        job is complete. If the output is a tty or a Jupyter cell, it will be color-coded
        based on which instance the log entry is from.

        Args:
            job_name (str): Name of the training job to display the logs for.
            wait (bool): Whether to keep looking for new log entries until the job completes (default: False).
            poll (int): The interval in seconds between polling for new log entries and job completion (default: 5).

        Raises:
            ValueError: If waiting and the training job fails.

def logs_for_job(self, job_name, wait=False, poll=10):  # noqa: C901 - suppress complexity warning for this method
        """Display the logs for a given training job, optionally tailing them until the
        job is complete. If the output is a tty or a Jupyter cell, it will be color-coded
        based on which instance the log entry is from.

        Args:
            job_name (str): Name of the training job to display the logs for.
            wait (bool): Whether to keep looking for new log entries until the job completes (default: False).
            poll (int): The interval in seconds between polling for new log entries and job completion (default: 5).

        Raises:
            ValueError: If waiting and the training job fails.
        """

        description = self.sagemaker_client.describe_training_job(TrainingJobName=job_name)
        print(secondary_training_status_message(description, None), end='')
        instance_count = description['ResourceConfig']['InstanceCount']
        status = description['TrainingJobStatus']

        stream_names = []  # The list of log streams
        positions = {}     # The current position in each stream, map of stream name -> position

        # Increase retries allowed (from default of 4), as we don't want waiting for a training job
        # to be interrupted by a transient exception.
        config = botocore.config.Config(retries={'max_attempts': 15})
        client = self.boto_session.client('logs', config=config)
        log_group = '/aws/sagemaker/TrainingJobs'

        job_already_completed = True if status == 'Completed' or status == 'Failed' or status == 'Stopped' else False

        state = LogState.TAILING if wait and not job_already_completed else LogState.COMPLETE
        dot = False

        color_wrap = sagemaker.logs.ColorWrap()

        # The loop below implements a state machine that alternates between checking the job status and
        # reading whatever is available in the logs at this point. Note, that if we were called with
        # wait == False, we never check the job status.
        #
        # If wait == TRUE and job is not completed, the initial state is TAILING
        # If wait == FALSE, the initial state is COMPLETE (doesn't matter if the job really is complete).
        #
        # The state table:
        #
        # STATE               ACTIONS                        CONDITION             NEW STATE
        # ----------------    ----------------               -----------------     ----------------
        # TAILING             Read logs, Pause, Get status   Job complete          JOB_COMPLETE
        #                                                    Else                  TAILING
        # JOB_COMPLETE        Read logs, Pause               Any                   COMPLETE
        # COMPLETE            Read logs, Exit                                      N/A
        #
        # Notes:
        # - The JOB_COMPLETE state forces us to do an extra pause and read any items that got to Cloudwatch after
        #   the job was marked complete.
        last_describe_job_call = time.time()
        last_description = description
        while True:
            if len(stream_names) < instance_count:
                # Log streams are created whenever a container starts writing to stdout/err, so this list
                # may be dynamic until we have a stream for every instance.
                try:
                    streams = client.describe_log_streams(logGroupName=log_group, logStreamNamePrefix=job_name + '/',
                                                          orderBy='LogStreamName', limit=instance_count)
                    stream_names = [s['logStreamName'] for s in streams['logStreams']]
                    positions.update([(s, sagemaker.logs.Position(timestamp=0, skip=0))
                                      for s in stream_names if s not in positions])
                except ClientError as e:
                    # On the very first training job run on an account, there's no log group until
                    # the container starts logging, so ignore any errors thrown about that
                    err = e.response.get('Error', {})
                    if err.get('Code', None) != 'ResourceNotFoundException':
                        raise

            if len(stream_names) > 0:
                if dot:
                    print('')
                    dot = False
                for idx, event in sagemaker.logs.multi_stream_iter(client, log_group, stream_names, positions):
                    color_wrap(idx, event['message'])
                    ts, count = positions[stream_names[idx]]
                    if event['timestamp'] == ts:
                        positions[stream_names[idx]] = sagemaker.logs.Position(timestamp=ts, skip=count + 1)
                    else:
                        positions[stream_names[idx]] = sagemaker.logs.Position(timestamp=event['timestamp'], skip=1)
            else:
                dot = True
                print('.', end='')
                sys.stdout.flush()
            if state == LogState.COMPLETE:
                break

            time.sleep(poll)

            if state == LogState.JOB_COMPLETE:
                state = LogState.COMPLETE
            elif time.time() - last_describe_job_call >= 30:
                description = self.sagemaker_client.describe_training_job(TrainingJobName=job_name)
                last_describe_job_call = time.time()

                if secondary_training_status_changed(description, last_description):
                    print()
                    print(secondary_training_status_message(description, last_description), end='')
                    last_description = description

                status = description['TrainingJobStatus']

                if status == 'Completed' or status == 'Failed' or status == 'Stopped':
                    print()
                    state = LogState.JOB_COMPLETE

        if wait:
            self._check_job_status(job_name, description, 'TrainingJobStatus')
            if dot:
                print()
            # Customers are not billed for hardware provisioning, so billable time is less than total time
            billable_time = (description['TrainingEndTime'] - description['TrainingStartTime']) * instance_count
            print('Billable seconds:', int(billable_time.total_seconds()) + 1)