# Copyright 2021 DeepMind Technologies Limited
#
# 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.
"""Local backend executors."""

from typing import Optional

import attr
from xmanager import xm


GOOGLE_KUBERNETES_ENGINE_CLOUD_PROVIDER = 'GOOGLE_KUBERNETES_ENGINE'


class LocalSpec(xm.ExecutorSpec):
  """Current machine executor's specification."""


class Local(xm.Executor):
  """Current machine executor."""

  @classmethod
  def Spec(cls):
    return LocalSpec()


@attr.s(auto_attribs=True)
class TpuCapability:
  """TPU capability configures the TPU software requested by an executor."""

  # Read about TPU versions:
  # https://cloud.google.com/tpu/docs/version-switching
  tpu_runtime_version: str


@attr.s(auto_attribs=True)
class CaipSpec(xm.ExecutorSpec):
  """Caip spec describes the Google Cloud Platform (GCP) location."""

  # An image registry name tag to push.
  # The image tag should be in the form 'myregistryhost/name:tag'
  push_image_tag: Optional[str] = None


@attr.s(auto_attribs=True)
class Caip(xm.Executor):
  """Caip Executor describes the runtime environment of GCP."""

  resources: xm.JobRequirements
  tpu_capability: Optional[TpuCapability] = None

  @classmethod
  def Spec(cls, *args, **kwargs):
    return CaipSpec(*args, **kwargs)


@attr.s(auto_attribs=True)
class KubernetesSpec(xm.ExecutorSpec):
  """K8s spec describes the K8s location."""

  # An image registry name tag to push.
  # The image tag should be in the form 'myregistryhost/name:tag'
  push_image_tag: Optional[str] = None


@attr.s(auto_attribs=True)
class Kubernetes(xm.Executor):
  """K8s Executor describes the runtime environment of Kubernetes."""

  resources: xm.JobRequirements
  cloud_provider: str = GOOGLE_KUBERNETES_ENGINE_CLOUD_PROVIDER
  tpu_capability: Optional[TpuCapability] = None

  @classmethod
  def Spec(cls, *args, **kwargs):
    return KubernetesSpec(*args, **kwargs)

@@ -638,24 +638,50 @@
  import 
-Optional
+Dict, List, Optional, Tuple, Union
 %0A%0Aimport
@@ -869,16 +869,773 @@
 n.%22%22%22%0A%0A%0A
+Ports = Dict%5BUnion%5Bint, str%5D, Union%5BNone, int, Tuple%5Bstr, int%5D, List%5Bint%5D%5D%5D%0A%0A%0A@attr.s(auto_attribs=True)%0Aclass DockerOptions:%0A  %22%22%22Options of the container to be run.%0A%0A  Attributes:%0A      ports: In the simplest form -- a dictionary from %60int%60 to %60int%60, where the%0A        keys represent the ports inside the container and the values represent%0A        the ports of the host to bind. See the specification at%0A        https://docker-py.readthedocs.io/en/stable/containers.html.%0A      volumes: A dictionary from %60str%60 to %60str%60, where the keys represent paths%0A        inside on the host to mount and the values represent paths in the%0A        container.%0A  %22%22%22%0A  ports: Optional%5BPorts%5D = None%0A  volumes: Optional%5BDict%5Bstr, str%5D%5D = None%0A%0A%0A@attr.s(auto_attribs=True)%0A
 class Lo
@@ -1687,16 +1687,66 @@
 or.%22%22%22%0A%0A
+  docker_options: Optional%5BDockerOptions%5D = None%0A%0A
   @class

Introduce `docker_options` to `xm_local.Local`