<commit_before># Copyright 2020 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Portable library for output artifacts resolution including caching decision."""

import collections
import datetime
import os
from typing import Dict, List, Optional, Text

from absl import logging
from tfx import types
from tfx import version
from tfx.dsl.io import fileio
from tfx.orchestration import data_types_utils
from tfx.proto.orchestration import execution_result_pb2
from tfx.proto.orchestration import pipeline_pb2
from tfx.types import artifact_utils
from tfx.types.value_artifact import ValueArtifact

_SYSTEM = '.system'
_EXECUTOR_EXECUTION = 'executor_execution'
_DRIVER_EXECUTION = 'driver_execution'
_STATEFUL_WORKING_DIR = 'stateful_working_dir'
_DRIVER_OUTPUT_FILE = 'driver_output.pb'
_EXECUTOR_OUTPUT_FILE = 'executor_output.pb'
_VALUE_ARTIFACT_FILE_NAME = 'value'
_ARTIFACT_TFX_VERSION_CUSTOM_PROPERTY_KEY = 'tfx_version'


def make_output_dirs(output_dict: Dict[Text, List[types.Artifact]]) -> None:
  """Make dirs for output artifacts' URI."""
  for _, artifact_list in output_dict.items():
    for artifact in artifact_list:
      if isinstance(artifact, ValueArtifact):
        # If it is a ValueArtifact, create a file.
        artifact_dir = os.path.dirname(artifact.uri)
        fileio.makedirs(artifact_dir)
        with fileio.open(artifact.uri, 'w') as f:
          # Because fileio.open won't create an empty file, we write an
          # empty string to it to force the creation.
          f.write('')
      else:
        # Otherwise create a dir.
        fileio.makedirs(artifact.uri)


def remove_output_dirs(output_dict: Dict[Text, List[types.Artifact]]) -> None:
  """Remove dirs of output artifacts' URI."""
  for _, artifact_list in output_dict.items():
    for artifact in artifact_list:
      if fileio.isdir(artifact.uri):
        fileio.rmtree(artifact.uri)
      else:
        fileio.remove(artifact.uri)


def remove_stateful_working_dir(stateful_working_dir: Text) -> None:
  """Remove stateful_working_dir."""
  # Clean up stateful working dir
  # Note that:
  # stateful_working_dir = os.path.join(
  #    self._node_dir,
  #    _SYSTEM,
  #    _STATEFUL_WORKING_DIR, <-- we want to clean from this level down.
  #    dir_suffix)
  stateful_working_dir = os.path.abspath(
      os.path.join(stateful_working_dir, os.pardir))
  try:
    fileio.rmtree(stateful_working_dir)
  except Exception as e:  # pylint: disable=broad-except
    if 'NotFoundError' in str(type(e)):
      # TODO(b/175244977): This is a workaround to avoid introducing
      # tensorflow dependency. Change this except block to use a generic
      # NotFoundError once it is Defined in fileio.
      logging.warning(
          'stateful_working_dir %s is not found, not going to delete it.',
          stateful_working_dir)
    else:
      raise


def _attach_artifact_properties(spec: pipeline_pb2.OutputSpec.ArtifactSpec,
                                artifact: types.Artifact):
  """Attaches properties of an artifact using ArtifactSpec."""
  for key, value in spec.additional_properties.items():
    if not value.HasField('field_value'):
      raise RuntimeError('Property value is not a field_value for %s' % key)<commit_msg>Addresses TODO as the generic NotFoundError is already available.<commit_after># Copyright 2020 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Portable library for output artifacts resolution including caching decision."""

import collections
import datetime
import os
from typing import Dict, List, Optional, Text

from absl import logging
from tfx import types
from tfx import version
from tfx.dsl.io import fileio
from tfx.orchestration import data_types_utils
from tfx.proto.orchestration import execution_result_pb2
from tfx.proto.orchestration import pipeline_pb2
from tfx.types import artifact_utils
from tfx.types.value_artifact import ValueArtifact

_SYSTEM = '.system'
_EXECUTOR_EXECUTION = 'executor_execution'
_DRIVER_EXECUTION = 'driver_execution'
_STATEFUL_WORKING_DIR = 'stateful_working_dir'
_DRIVER_OUTPUT_FILE = 'driver_output.pb'
_EXECUTOR_OUTPUT_FILE = 'executor_output.pb'
_VALUE_ARTIFACT_FILE_NAME = 'value'
_ARTIFACT_TFX_VERSION_CUSTOM_PROPERTY_KEY = 'tfx_version'


def make_output_dirs(output_dict: Dict[Text, List[types.Artifact]]) -> None:
  """Make dirs for output artifacts' URI."""
  for _, artifact_list in output_dict.items():
    for artifact in artifact_list:
      if isinstance(artifact, ValueArtifact):
        # If it is a ValueArtifact, create a file.
        artifact_dir = os.path.dirname(artifact.uri)
        fileio.makedirs(artifact_dir)
        with fileio.open(artifact.uri, 'w') as f:
          # Because fileio.open won't create an empty file, we write an
          # empty string to it to force the creation.
          f.write('')
      else:
        # Otherwise create a dir.
        fileio.makedirs(artifact.uri)


def remove_output_dirs(output_dict: Dict[Text, List[types.Artifact]]) -> None:
  """Remove dirs of output artifacts' URI."""
  for _, artifact_list in output_dict.items():
    for artifact in artifact_list:
      if fileio.isdir(artifact.uri):
        fileio.rmtree(artifact.uri)
      else:
        fileio.remove(artifact.uri)


def remove_stateful_working_dir(stateful_working_dir: Text) -> None:
  """Remove stateful_working_dir."""
  # Clean up stateful working dir
  # Note that:
  # stateful_working_dir = os.path.join(
  #    self._node_dir,
  #    _SYSTEM,
  #    _STATEFUL_WORKING_DIR, <-- we want to clean from this level down.
  #    dir_suffix)
  stateful_working_dir = os.path.abspath(
      os.path.join(stateful_working_dir, os.pardir))
  try:
    fileio.rmtree(stateful_working_dir)
  except fileio.NotFoundError:
    logging.warning(
        'stateful_working_dir %s is not found, not going to delete it.',
        stateful_working_dir)


def _attach_artifact_properties(spec: pipeline_pb2.OutputSpec.ArtifactSpec,
                                artifact: types.Artifact):
  """Attaches properties of an artifact using ArtifactSpec."""
  for key, value in spec.additional_properties.items():
    if not value.HasField('field_value'):
      raise RuntimeError('Property value is not a field_value for %s' % key)
{'fraction_non_alphanumeric': 0.08267068823690704, 'fraction_numerical': 0.0042500685494927335, 'mean_word_length': 3.879598662207358, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '4906001', 'n_tokens_mistral': 2247, 'n_tokens_neox': 2164, 'n_words': 690}