feat: Add new gcloud commands, API clients, and third-party libraries across various services.

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/command_lib/ai/docker/build.py

@@ -0,0 +1,338 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2021 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.
+"""Functions required to interact with Docker to build images."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+import json
+import os
+import posixpath
+import re
+import textwrap
+
+from googlecloudsdk.command_lib.ai import errors
+from googlecloudsdk.command_lib.ai.custom_jobs import local_util
+from googlecloudsdk.command_lib.ai.docker import utils
+from googlecloudsdk.core import log
+from six.moves import shlex_quote
+
+_DEFAULT_HOME = "/home"
+_DEFAULT_WORKDIR = "/usr/app"
+_DEFAULT_SETUP_PATH = "./setup.py"
+_DEFAULT_REQUIREMENTS_PATH = "./requirements.txt"
+_AUTONAME_PREFIX = "cloudai-autogenerated"
+_AUTOGENERATED_TAG_LENGTH = 16
+
+
+def _IsVertexTrainingPrebuiltImage(image_name):
+  """Checks whether the image is pre-built by Vertex AI training."""
+  prebuilt_image_name_regex = (r"^(us|europe|asia)-docker.pkg.dev/"
+                               r"vertex-ai/training/"
+                               r"(tf|scikit-learn|pytorch|xgboost)-.+$")
+  return re.fullmatch(prebuilt_image_name_regex, image_name) is not None
+
+
+def _SitecustomizeRemovalEntry(is_prebuilt_image):
+  """Returns a Dockerfile entry that removes `sitecustomize` if it's Vertex AI Training pre-built container images."""
+  return "RUN rm -rf /var/sitecustomize" if is_prebuilt_image else ""
+
+
+def _GenerateCopyCommand(from_path, to_path, comment=None):
+  """Returns a Dockerfile entry that copies a file from host to container.
+
+  Args:
+    from_path: (str) Path of the source in host.
+    to_path: (str) Path to the destination in the container.
+    comment: (str) A comment explaining the copy operation.
+  """
+  cmd = "COPY {}\n".format(json.dumps([from_path, to_path]))
+
+  if comment is not None:
+    formatted_comment = "\n# ".join(comment.split("\n"))
+    return "# {}\n{}".format(formatted_comment, cmd)
+
+  return cmd
+
+
+def _DependencyEntries(is_prebuilt_image=False,
+                       requirements_path=None,
+                       setup_path=None,
+                       extra_requirements=None,
+                       extra_packages=None,
+                       extra_dirs=None):
+  """Returns the Dockerfile entries required to install dependencies.
+
+  Args:
+    is_prebuilt_image: (bool) Whether the base image is pre-built and provided
+      by Vertex AI.
+    requirements_path: (str) Path that points to a requirements.txt file
+    setup_path: (str) Path that points to a setup.py
+    extra_requirements: (List[str]) Required dependencies to be installed from
+      remote resource archives.
+    extra_packages: (List[str]) User custom dependency packages to install.
+    extra_dirs: (List[str]) Directories other than the work_dir required.
+  """
+  ret = ""
+  pip_version = "pip3" if is_prebuilt_image else "pip"
+
+  if setup_path is not None:
+    ret += textwrap.dedent("""
+        {}
+        RUN {} install --no-cache-dir .
+        """.format(
+            _GenerateCopyCommand(
+                setup_path,
+                "./setup.py",
+                comment="Found setup.py file, thus copy it to the docker container."
+            ), pip_version))
+
+  if requirements_path is not None:
+    ret += textwrap.dedent("""
+        {}
+        RUN {} install --no-cache-dir -r ./requirements.txt
+        """.format(
+            _GenerateCopyCommand(
+                requirements_path,
+                "./requirements.txt",
+                comment="Found requirements.txt file, thus to the docker container."
+            ), pip_version))
+
+  if extra_packages is not None:
+    for extra in extra_packages:
+      package_name = os.path.basename(extra)
+      ret += textwrap.dedent("""
+        {}
+        RUN {} install --no-cache-dir {}
+        """.format(
+            _GenerateCopyCommand(extra, package_name), pip_version,
+            shlex_quote(package_name)))
+
+  if extra_requirements is not None:
+    for requirement in extra_requirements:
+      ret += textwrap.dedent("""
+        RUN {} install --no-cache-dir --upgrade {}
+        """.format(pip_version, shlex_quote(requirement)))
+
+  if extra_dirs is not None:
+    for directory in extra_dirs:
+      ret += "\n{}\n".format(_GenerateCopyCommand(directory, directory))
+
+  return ret
+
+
+def _GenerateEntrypoint(package, is_prebuilt_image=False):
+  """Generates dockerfile entry to set the container entrypoint.
+
+  Args:
+    package: (Package) Represents the main application copied to the container.
+    is_prebuilt_image: (bool) Whether the base image is pre-built and provided
+      by Vertex AI.
+
+  Returns:
+    A string with Dockerfile directives to set ENTRYPOINT
+  """
+
+  # Make it consistent with Online python package training that python3
+  # has been installed for all prebuilt images and used by default
+  python_command = "python3" if is_prebuilt_image else "python"
+
+  # Needs to use json so that quotes print as double quotes, not single quotes.
+  if package.python_module is not None:
+    exec_str = json.dumps([python_command, "-m", package.python_module])
+  else:
+    _, ext = os.path.splitext(package.script)
+    executable = [python_command] if ext == ".py" else ["/bin/bash"]
+    exec_str = json.dumps(executable + [package.script])
+
+  return "\nENTRYPOINT {}".format(exec_str)
+
+
+def _PreparePackageEntry(package):
+  """Returns the Dockerfile entries required to append at the end before entrypoint.
+
+  Including:
+  - copy the parent directory of the main executable into a docker container.
+  - inject an entrypoint that executes a script or python module inside that
+    directory.
+
+  Args:
+    package: (Package) Represents the main application copied to and run in the
+      container.
+  """
+  parent_dir = os.path.dirname(package.script) or "."
+
+  copy_code = _GenerateCopyCommand(
+      parent_dir,
+      parent_dir,
+      comment="Copy the source directory into the docker container.")
+
+  return "\n{}\n".format(copy_code)
+
+
+def _MakeDockerfile(base_image,
+                    main_package,
+                    container_workdir,
+                    container_home,
+                    requirements_path=None,
+                    setup_path=None,
+                    extra_requirements=None,
+                    extra_packages=None,
+                    extra_dirs=None):
+  """Generates a Dockerfile for building an image.
+
+  It builds on a specified base image to create a container that:
+  - installs any dependency specified in a requirements.txt or a setup.py file,
+  and any specified dependency packages existing locally or found from PyPI
+  - copies all source needed by the main module, and potentially injects an
+  entrypoint that, on run, will run that main module
+
+  Args:
+    base_image: (str) ID or name of the base image to initialize the build
+      stage.
+    main_package: (Package) Represents the main application to execute.
+    container_workdir: (str) Working directory in the container.
+    container_home: (str) $HOME directory in the container.
+    requirements_path: (str) Rath of a requirements.txt file.
+    setup_path: (str) Path of a setup.py file
+    extra_requirements: (List[str]) Required dependencies to install from PyPI.
+    extra_packages: (List[str]) User custom dependency packages to install.
+    extra_dirs: (List[str]) Directories other than the work_dir required to be
+      in the container.
+
+  Returns:
+    A string that represents the content of a Dockerfile.
+  """
+  is_training_prebuilt_image_base = _IsVertexTrainingPrebuiltImage(base_image)
+
+  dockerfile = textwrap.dedent("""
+      FROM {base_image}
+      # The directory is created by root. This sets permissions so that any user can
+      # access the folder.
+      RUN mkdir -m 777 -p {workdir} {container_home}
+      WORKDIR {workdir}
+      ENV HOME={container_home}
+
+      # Keeps Python from generating .pyc files in the container
+      ENV PYTHONDONTWRITEBYTECODE=1
+      """.format(
+          base_image=base_image,
+          workdir=shlex_quote(container_workdir),
+          container_home=shlex_quote(container_home)))
+
+  dockerfile += _SitecustomizeRemovalEntry(is_training_prebuilt_image_base)
+
+  dockerfile += _DependencyEntries(
+      is_training_prebuilt_image_base,
+      requirements_path=requirements_path,
+      setup_path=setup_path,
+      extra_requirements=extra_requirements,
+      extra_packages=extra_packages,
+      extra_dirs=extra_dirs)
+
+  dockerfile += _PreparePackageEntry(main_package)
+  dockerfile += _GenerateEntrypoint(main_package,
+                                    is_training_prebuilt_image_base)
+
+  return dockerfile
+
+
+def BuildImage(base_image,
+               host_workdir,
+               main_script,
+               output_image_name,
+               python_module=None,
+               requirements=None,
+               extra_packages=None,
+               container_workdir=None,
+               container_home=None,
+               no_cache=True,
+               **kwargs):
+  """Builds a Docker image.
+
+  Generates a Dockerfile and passes it to `docker build` via stdin.
+  All output from the `docker build` process prints to stdout.
+
+  Args:
+    base_image: (str) ID or name of the base image to initialize the build
+      stage.
+    host_workdir: (str) A path indicating where all the required sources
+      locates.
+    main_script: (str) A string that identifies the executable script under the
+      working directory.
+    output_image_name: (str) Name of the built image.
+    python_module: (str) Represents the executable main_script in form of a
+      python module, if applicable.
+    requirements: (List[str]) Required dependencies to install from PyPI.
+    extra_packages: (List[str]) User custom dependency packages to install.
+    container_workdir: (str) Working directory in the container.
+    container_home: (str) the $HOME directory in the container.
+    no_cache: (bool) Do not use cache when building the image.
+    **kwargs: Other arguments to pass to underlying method that generates the
+      Dockerfile.
+
+  Returns:
+    A Image class that contains info of the built image.
+
+  Raises:
+    DockerError: An error occurred when executing `docker build`
+  """
+
+  tag_options = ["-t", output_image_name]
+
+  cache_args = ["--no-cache"] if no_cache else []
+  command = ["docker", "build"
+            ] + cache_args + tag_options + ["--rm", "-f-", host_workdir]
+
+  has_setup_py = os.path.isfile(os.path.join(host_workdir, _DEFAULT_SETUP_PATH))
+  setup_path = _DEFAULT_SETUP_PATH if has_setup_py else None
+
+  has_requirements_txt = os.path.isfile(
+      os.path.join(host_workdir, _DEFAULT_REQUIREMENTS_PATH))
+  requirements_path = _DEFAULT_REQUIREMENTS_PATH if has_requirements_txt else None
+
+  home_dir = container_home or _DEFAULT_HOME
+  work_dir = container_workdir or _DEFAULT_WORKDIR
+
+  # The package will be used in Docker, thus norm it to POSIX path format.
+  main_package = utils.Package(
+      script=main_script.replace(os.sep, posixpath.sep),
+      package_path=host_workdir.replace(os.sep, posixpath.sep),
+      python_module=python_module)
+
+  dockerfile = _MakeDockerfile(
+      base_image,
+      main_package=main_package,
+      container_home=home_dir,
+      container_workdir=work_dir,
+      requirements_path=requirements_path,
+      setup_path=setup_path,
+      extra_requirements=requirements,
+      extra_packages=extra_packages,
+      **kwargs)
+
+  joined_command = " ".join(command)
+  log.info("Running command: {}".format(joined_command))
+
+  return_code = local_util.ExecuteCommand(command, input_str=dockerfile)
+  if return_code == 0:
+    return utils.Image(output_image_name, home_dir, work_dir)
+  else:
+    error_msg = textwrap.dedent("""
+        Docker failed with error code {code}.
+        Command: {cmd}
+        """.format(code=return_code, cmd=joined_command))
+    raise errors.DockerError(error_msg, command, return_code)

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/command_lib/ai/docker/run.py

@@ -0,0 +1,87 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2021 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.
+"""Functions required to interact with Docker to run a container."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+from googlecloudsdk.command_lib.ai.docker import utils
+from googlecloudsdk.core import config
+
+_DEFAULT_CONTAINER_CRED_KEY_PATH = "/tmp/keys/cred_key.json"
+
+
+def _DockerRunOptions(enable_gpu=False,
+                      service_account_key=None,
+                      cred_mount_path=_DEFAULT_CONTAINER_CRED_KEY_PATH,
+                      extra_run_opts=None):
+  """Returns a list of 'docker run' options.
+
+  Args:
+    enable_gpu: (bool) using GPU or not.
+    service_account_key: (bool) path of the service account key to use in host.
+    cred_mount_path: (str) path in the container to mount the credential key.
+    extra_run_opts: (List[str]) other custom docker run options.
+  """
+  if extra_run_opts is None:
+    extra_run_opts = []
+
+  runtime = ["--runtime", "nvidia"] if enable_gpu else []
+
+  if service_account_key:
+    mount = ["-v", "{}:{}".format(service_account_key, cred_mount_path)]
+  else:
+    # Calls Application Default Credential (ADC),
+    adc_file_path = config.ADCEnvVariable() or config.ADCFilePath()
+    mount = ["-v", "{}:{}".format(adc_file_path, cred_mount_path)]
+  env_var = ["-e", "GOOGLE_APPLICATION_CREDENTIALS={}".format(cred_mount_path)]
+
+  return ["--rm"] + runtime + mount + env_var + ["--ipc", "host"
+                                                ] + extra_run_opts
+
+
+def RunContainer(image_name,
+                 enable_gpu=False,
+                 service_account_key=None,
+                 run_args=None,
+                 user_args=None):
+  """Calls `docker run` on a given image with specified arguments.
+
+  Args:
+    image_name: (str) Name or ID of Docker image to run.
+    enable_gpu: (bool) Whether to use GPU
+    service_account_key: (str) Json file of a service account key  auth.
+    run_args: (List[str]) Extra custom options to apply to `docker run` after
+      our defaults.
+    user_args: (List[str]) Extra user defined arguments to supply to the
+      entrypoint.
+  """
+  # TODO(b/177787660): add interactive mode option
+
+  if run_args is None:
+    run_args = []
+
+  if user_args is None:
+    user_args = []
+
+  run_opts = _DockerRunOptions(
+      enable_gpu=enable_gpu,
+      service_account_key=service_account_key,
+      extra_run_opts=run_args)
+
+  command = ["docker", "run"] + run_opts + [image_name] + user_args
+
+  utils.ExecuteDockerCommand(command)

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/command_lib/ai/docker/utils.py

@@ -0,0 +1,244 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2021 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.
+"""Common utilities to operate with Docker."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+import collections
+import datetime
+import re
+from googlecloudsdk.command_lib.ai import errors
+from googlecloudsdk.command_lib.ai.custom_jobs import local_util
+from googlecloudsdk.core import log
+
+_MAX_REPOSITORY_LENGTH = 255
+_MAX_TAG_LENGTH = 128
+
+_AUTONAME_PREFIX = "cloudai-autogenerated"
+_DEFAULT_IMAGE_NAME = "unnamed"
+_DEFAULT_REPO_REGION = "us"
+
+Package = collections.namedtuple("Package",
+                                 ["script", "package_path", "python_module"])
+Image = collections.namedtuple("Image",
+                               ["name", "default_home", "default_workdir"])
+
+
+def _ParseRepositoryTag(image_name):
+  """Parses out the repository and tag from a Docker image name.
+
+  Args:
+    image_name: (str) The full name of an image, expected to be in a format of
+      "repository[:tag]"
+
+  Returns:
+    A (repository, tag) tuple representing the parsed result.
+    None repository means the image name is invalid; tag may be None if it isn't
+    present in the given image name.
+  """
+
+  if image_name.count(":") > 2:
+    return None, None
+
+  parts = image_name.rsplit(":", 1)
+  if len(parts) == 2 and "/" not in parts[1]:
+    return tuple(parts)
+
+  return image_name, None
+
+
+def _ParseRepositoryHost(repository_name):
+  """Parses a repository to an optional hostname and a list of path compoentes.
+
+  Args:
+    repository_name: (str) A name made up of slash-separated path name
+      components, optionally prefixed by a registry hostname.
+
+  Returns:
+    A (hostname, components) tuple representing the parsed result.
+    The hostname will be None if it isn't present; the components is a list of
+    each slash-separated part in the given repository name.
+  """
+
+  components = repository_name.split("/")
+
+  if len(components) == 1:
+    return None, components
+
+  if "." in components[0] or ":" in components[0]:
+    # components[0] is regarded as a hostname
+    return components[0], components[1:]
+
+  return None, components
+
+
+def _ParseHostPort(host):
+  """Parses a registry hostname to a list of components and an optional port.
+
+  Args:
+    host: (str) The registry hostname supposed to comply with standard DNS
+      rules, optionally be followed by a port number in the format like ":8080".
+
+  Returns:
+    A (hostcomponents, port) tuple representing the parsed result.
+    The hostcomponents contains each dot-seperated component in the given
+    hostname; port may be None if it isn't present.
+  """
+
+  parts = host.rsplit(":", 1)
+  hostcomponents = parts[0].split(".")
+  port = parts[1] if len(parts) == 2 else None
+
+  return hostcomponents, port
+
+
+def ValidateRepositoryAndTag(image_name):
+  r"""Validate the given image name is a valid repository/tag reference.
+
+  As explained in
+  https://docs.docker.com/engine/reference/commandline/tag/#extended-description,
+  a valid repository/tag reference should following the below pattern:
+
+  reference             := name [ ":" tag ]
+  name                  := [hostname '/'] component ['/' component]*
+  hostname              := hostcomponent ['.' hostcomponent]* [':' port-number]
+  hostcomponent         := /([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9-]*[a-zA-Z0-9])/
+  port-number           := /[0-9]+/
+  component             := alpha-numeric [separator alpha-numeric]*
+  alpha-numeric         := /[a-z0-9]+/
+  separator             := /[_.]|__|[-]*/
+
+  tag                   := /[\w][\w.-]{0,127}/
+
+  Args:
+    image_name: (str) Full name of a Docker image.
+
+  Raises:
+    ValueError if the image name is not valid.
+  """
+  repository, tag = _ParseRepositoryTag(image_name)
+
+  if repository is None:
+    raise ValueError("Unable to parse repository and tag.")
+
+  if len(repository) > _MAX_REPOSITORY_LENGTH:
+    raise ValueError(
+        "Repository name must not be more than {} characters.".format(
+            _MAX_REPOSITORY_LENGTH))
+
+  hostname, path_components = _ParseRepositoryHost(repository)
+
+  if hostname:
+    hostcomponents, port = _ParseHostPort(hostname)
+
+    hostcomponent_regex = r"^(?:[a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9-]*[a-zA-Z0-9])$"
+    for hostcomponent in hostcomponents:
+      if re.match(hostcomponent_regex, hostcomponent) is None:
+        raise ValueError(
+            "Invalid hostname/port \"{}\" in repository name.".format(hostname))
+
+    port_regex = r"^[0-9]+$"
+    if port and re.match(port_regex, port) is None:
+      raise ValueError(
+          "Invalid hostname/port \"{}\" in repository name.".format(hostname))
+
+  for component in path_components:
+    if not component:
+      raise ValueError("Empty path component in repository name.")
+
+    component_regex = r"^[a-z0-9]+(?:(?:[._]|__|[-]*)[a-z0-9]+)*$"
+    if re.match(component_regex, component) is None:
+      raise ValueError(
+          "Invalid path component \"{}\" in repository name.".format(component))
+
+  if tag:
+    if len(tag) > _MAX_TAG_LENGTH:
+      raise ValueError("Tag name must not be more than {} characters.".format(
+          _MAX_TAG_LENGTH))
+
+    tag_regex = r"^[\w][\w.-]{0,127}$"
+    if re.match(tag_regex, tag) is None:
+      raise ValueError("Invalid tag.")
+
+
+def GenerateImageName(base_name=None, project=None, region=None, is_gcr=False):
+  """Generate a name for the Docker image built by AI platform gcloud."""
+  sanitized_name = _SanitizeRepositoryName(base_name or _DEFAULT_IMAGE_NAME)
+
+  # Use the current timestamp as the tag.
+  tag = datetime.datetime.now().strftime("%Y%m%d.%H.%M.%S.%f")
+  image_name = "{}/{}:{}".format(_AUTONAME_PREFIX, sanitized_name, tag)
+
+  if project:
+    if is_gcr:
+      repository = "gcr.io"
+    else:
+      region_prefix = region or _DEFAULT_REPO_REGION
+      repository = "{}-docker.pkg.dev".format(region_prefix)
+    return "{}/{}/{}".format(repository, project.replace(":", "/"), image_name)
+  return image_name
+
+
+def _SanitizeRepositoryName(name):
+  """Sanitizes the given name to make it valid as an image repository.
+
+  As explained in
+  https://docs.docker.com/engine/reference/commandline/tag/#extended-description,
+  Valid name may contain only lowercase letters, digits and separators.
+  A separator is defined as a period, one or two underscores, or one or more
+  dashes. A name component may not start or end with a separator.
+
+  This method will replace the illegal characters in the given name and strip
+  starting and ending separator characters.
+
+  Args:
+    name: str, the name to sanitize.
+
+  Returns:
+    A sanitized name.
+  """
+  return re.sub("[._][._]+|[^a-z0-9._-]+", ".", name.lower()).strip("._-")
+
+
+def ExecuteDockerCommand(command):
+  """Executes Docker CLI commands in subprocess.
+
+  Just calls local_util.ExecuteCommand(cmd,...) and raises error for non-zero
+  exit code.
+
+  Args:
+    command: (List[str]) Strings to send in as the command.
+
+  Raises:
+    ValueError: The input command is not a docker command.
+    DockerError: An error occurred when executing the given docker command.
+  """
+
+  command_str = " ".join(command)
+  if not command_str.startswith("docker"):
+    raise ValueError("`{}` is not a Docker command".format("docker"))
+
+  log.info("Running command: {}".format(command_str))
+
+  return_code = local_util.ExecuteCommand(command)
+  if return_code != 0:
+    error_msg = """
+        Docker failed with error code {code}.
+        Command: {cmd}
+        """.format(
+            code=return_code, cmd=command_str)
+    raise errors.DockerError(error_msg, command, return_code)