class_docstr: Interact with Amazon SageMaker.
  
  Provide thick wrapper around
  :external+boto3:py:class:`boto3.client("sagemaker") <SageMaker.Client>`.
  
  Additional arguments (such as ``aws_conn_id``) may be specified and
  are passed down to the underlying AwsBaseHook.
  
  .. seealso::
      - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
class_docstr_tok: ["Interact","with","Amazon","SageMaker",".","Provide","thick","wrapper","around",":","external+boto3",":","py",":","class",":","`","boto3.client","(","``","sagemaker","''",")","<","SageMaker.Client",">","`",".","Additional","arguments","(","such","as","``","aws_conn_id","``",")","may","be","specified","and","are","passed","down","to","the","underlying","AwsBaseHook",".","..","seealso",":",":","-",":","class",":","`","airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook","`"]
docstr: Display logs for a given training job.
  
  Optionally tailing them until the job is complete.
  
  :param job_name: name of the training job to check status and display logs for
  :param non_terminal_states: the set of non_terminal states
  :param failed_states: the set of failed states
  :param wait_for_completion: Whether to keep looking for new log entries
      until the job completes
  :param check_interval: The interval in seconds between polling for new log entries and job completion
  :param max_ingestion_time: the maximum ingestion time in seconds. Any
      SageMaker jobs that run longer than this will fail. Setting this to
      None implies no timeout for any SageMaker job.
code: def check_training_status_with_log(
      self,
      job_name: str,
      non_terminal_states: set,
      failed_states: set,
      wait_for_completion: bool,
      check_interval: int,
      max_ingestion_time: int | None = None,
  ):
      """Display logs for a given training job.
  
      Optionally tailing them until the job is complete.
  
      :param job_name: name of the training job to check status and display logs for
      :param non_terminal_states: the set of non_terminal states
      :param failed_states: the set of failed states
      :param wait_for_completion: Whether to keep looking for new log entries
          until the job completes
      :param check_interval: The interval in seconds between polling for new log entries and job completion
      :param max_ingestion_time: the maximum ingestion time in seconds. Any
          SageMaker jobs that run longer than this will fail. Setting this to
          None implies no timeout for any SageMaker job.
      """
      sec = 0
      description = self.describe_training_job(job_name)
      self.log.info(
          secondary_training_status_message(description, None)
      )
      instance_count = description["ResourceConfig"]["InstanceCount"]
      status = description["TrainingJobStatus"]
  
      stream_names: list = []  # The list of log streams
      positions: dict = (
          {}
      )  # The current position in each stream, map of stream name -> position
  
      job_already_completed = status not in non_terminal_states
  
      state = (
          LogState.TAILING
          if wait_for_completion and not job_already_completed
          else LogState.COMPLETE
      )
  
      # 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_for_completion == False, we never check the job status.
      #
      # If wait_for_completion == TRUE and job is not completed, the initial state is TAILING
      # If wait_for_completion == 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.monotonic()
      last_description = description
  
      while True:
          time.sleep(check_interval)
          sec += check_interval
  
          (
              state,
              last_description,
              last_describe_job_call,
          ) = self.describe_training_job_with_log(
              job_name,
              positions,
              stream_names,
              instance_count,
              state,
              last_description,
              last_describe_job_call,
          )
          if state == LogState.COMPLETE:
              break
  
          if max_ingestion_time and sec > max_ingestion_time:
              # ensure that the job gets killed if the max ingestion time is exceeded
              raise AirflowException(
                  f"SageMaker job took more than {max_ingestion_time} seconds"
              )
  
      if wait_for_completion:
          status = last_description["TrainingJobStatus"]
          if status in failed_states:
              reason = last_description.get(
                  "FailureReason", "(No reason provided)"
              )
              raise AirflowException(
                  f"Error training {job_name}: {status} Reason: {reason}"
              )
          billable_time = (
              last_description["TrainingEndTime"]
              - last_description["TrainingStartTime"]
          ) * instance_count
          self.log.info(
              "Billable seconds: %d",
              int(billable_time.total_seconds()) + 1,
          )
  
imports: from __future__ import annotations
  import os
  import re
  import tarfile
  import tempfile
  import time
  import warnings
  from collections import Counter, namedtuple
  from datetime import datetime
  from functools import partial
  from typing import Any, Callable, Generator, cast
  from botocore.exceptions import ClientError
  from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
  from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
  from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
  from airflow.providers.amazon.aws.hooks.s3 import S3Hook
  from airflow.providers.amazon.aws.utils.tags import format_tags
  from airflow.utils import timezone
label_desc: Use image node_id 25 for calling the SageMakerHook obj's underlying member method code with example usage: obj.check_training_status_with_log(job_name, non_terminal_states, failed_states, wait_for_completion, check_interval, max_ingestion_time) without return types