<commit_before>"""Some utility functions for working with TfJobs."""

import datetime
import logging
from retrying import retry
import time

from kubernetes import client as k8s_client
from kubernetes.client import rest

from kubeflow.testing import util

GROUP = "argoproj.io"
VERSION = "v1alpha1"
PLURAL = "workflows"
KIND = "Workflow"

def log_status(workflow):
  """A callback to use with wait_for_workflow."""
  try:
    logging.info("Workflow %s in namespace %s; phase=%s",
                 workflow["metadata"]["name"],
                 workflow["metadata"]["namespace"],
                 workflow["status"]["phase"])
  except KeyError as e:
    # Ignore the error and just log the stacktrace
    # as sometimes the workflow object does not have all the fields
    # https://github.com/kubeflow/testing/issues/147
    logging.exception('KeyError: %s', e)


def handle_retriable_exception(exception):
  if (isinstance(exception, rest.ApiException) and
    (exception.status == 401 or exception.status == 403)):
    # Due to https://github.com/kubernetes-client/python-base/issues/59,
    # we need to reload the kube config (which refreshes the GCP token).
    # TODO(richardsliu): Remove this workaround when the k8s client issue
    # is resolved.
    util.load_kube_config()<commit_msg>Retry on 504 GATEWAY_TIMEOUT (#395)<commit_after>"""Some utility functions for working with TfJobs."""

import datetime
import logging
from retrying import retry
import time

from kubernetes import client as k8s_client
from kubernetes.client import rest

from kubeflow.testing import util

GROUP = "argoproj.io"
VERSION = "v1alpha1"
PLURAL = "workflows"
KIND = "Workflow"

def log_status(workflow):
  """A callback to use with wait_for_workflow."""
  try:
    logging.info("Workflow %s in namespace %s; phase=%s",
                 workflow["metadata"]["name"],
                 workflow["metadata"]["namespace"],
                 workflow["status"]["phase"])
  except KeyError as e:
    # Ignore the error and just log the stacktrace
    # as sometimes the workflow object does not have all the fields
    # https://github.com/kubeflow/testing/issues/147
    logging.exception('KeyError: %s', e)

# TODO(jlewi):
# In python3 we can switch to using http.HttpStatusCode
UNAUTHORIZED = 401
FORBIDDEN = 403
GATEWAY_TIMEOUT = 504

def handle_retriable_exception(exception):
  if (isinstance(exception, rest.ApiException) and
    # UNAUTHORIZED and FORBIDDEN errors can be an indication we need to
    # refresh credentials
    (exception.status == UNAUTHORIZED or exception.status == FORBIDDEN or
     exception.status == GATEWAY_TIMEOUT)):
    # Due to https://github.com/kubernetes-client/python-base/issues/59,
    # we need to reload the kube config (which refreshes the GCP token).
    # TODO(richardsliu): Remove this workaround when the k8s client issue
    # is resolved.
    util.load_kube_config()
{'fraction_non_alphanumeric': 0.09975837072833967, 'fraction_numerical': 0.013807386952019331, 'mean_word_length': 3.7900826446280993, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4919288', 'n_tokens_mistral': 891, 'n_tokens_neox': 865, 'n_words': 278}