From 48f2c35cee7715da996408845a8b6dab745da1d8 Mon Sep 17 00:00:00 2001
From: DeepMind Team <no-reply@deepmind.com>
Date: Wed, 31 Jul 2024 16:43:52 -0700
Subject: [PATCH] Don't restrict values for env_vars to str.

PiperOrigin-RevId: 658186331
Change-Id: I85504c2e561729687ab985887815fb72c5c9d608
GitOrigin-RevId: f07d2926d4856cfe69bd6e5d3f58d147920a74e0
---
 xmanager/xm/job_blocks.py  |  6 +++---
 xmanager/xm/packagables.py | 14 +++++++-------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/xmanager/xm/job_blocks.py b/xmanager/xm/job_blocks.py
index dfa691a..80056f4 100644
--- a/xmanager/xm/job_blocks.py
+++ b/xmanager/xm/job_blocks.py
@@ -341,7 +341,7 @@ async def launch(
 
 
 def _validate_env_vars(
-    self: Any, attribute: Any, env_vars: Dict[str, str]
+    self: Any, attribute: Any, env_vars: Dict[str, Any]
 ) -> None:
   del self  # Unused.
   del attribute  # Unused.
@@ -362,7 +362,7 @@ class Packageable:
   args: SequentialArgs = attr.ib(
       factory=list, converter=SequentialArgs.from_collection
   )  # pytype: disable=annotation-type-mismatch
-  env_vars: Dict[str, str] = attr.ib(
+  env_vars: Dict[str, Any] = attr.ib(
       converter=dict, default=attr.Factory(dict), validator=_validate_env_vars
   )
 
@@ -409,7 +409,7 @@ class Job:
   args: SequentialArgs = attr.ib(
       factory=list, converter=SequentialArgs.from_collection
   )  # pytype: disable=annotation-type-mismatch
-  env_vars: Dict[str, str] = attr.ib(
+  env_vars: Dict[str, Any] = attr.ib(
       converter=dict, default=attr.Factory(dict), validator=_validate_env_vars
   )
 
diff --git a/xmanager/xm/packagables.py b/xmanager/xm/packagables.py
index a282a9c..38a3ec4 100644
--- a/xmanager/xm/packagables.py
+++ b/xmanager/xm/packagables.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """Convenience methods for constructing core objects."""
 
-from typing import Collection, List, Mapping, Optional, Union
+from typing import Any, Collection, List, Mapping, Optional, Union
 
 import immutabledict
 from xmanager.xm import executables
@@ -30,7 +30,7 @@ def binary(
     dependencies: Collection[executables.BinaryDependency] = (),
     *,
     args: Optional[job_blocks.UserArgs] = None,
-    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+    env_vars: Mapping[str, Any] = immutabledict.immutabledict(),
 ) -> job_blocks.Packageable:
   # pyformat: disable
   """A prebuilt executable program.
@@ -70,7 +70,7 @@ def bazel_binary(
     bazel_args: Collection[str] = (),
     *,
     args: Optional[job_blocks.UserArgs] = None,
-    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+    env_vars: Mapping[str, Any] = immutabledict.immutabledict(),
 ) -> job_blocks.Packageable:
   # pyformat: disable
   """A Bazel target that produces a self-contained binary.
@@ -113,7 +113,7 @@ def container(
     image_path: str,
     *,
     args: Optional[job_blocks.UserArgs] = None,
-    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+    env_vars: Mapping[str, Any] = immutabledict.immutabledict(),
 ) -> job_blocks.Packageable:
   # pyformat: disable
   """A prebuilt Docker image.
@@ -151,7 +151,7 @@ def bazel_container(
     bazel_args: Collection[str] = (),
     *,
     args: Optional[job_blocks.UserArgs] = None,
-    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+    env_vars: Mapping[str, Any] = immutabledict.immutabledict(),
 ) -> job_blocks.Packageable:
   # pyformat: disable
   """A Bazel target that produces a .tar image.
@@ -195,7 +195,7 @@ def python_container(
     use_deep_module: bool = False,
     *,
     args: Optional[job_blocks.UserArgs] = None,
-    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+    env_vars: Mapping[str, Any] = immutabledict.immutabledict(),
 ) -> job_blocks.Packageable:
   # pyformat: disable
   """PythonContainer describes a directory containing Python code.
@@ -271,7 +271,7 @@ def dockerfile_container(
     dockerfile: Optional[str] = None,
     *,
     args: Optional[job_blocks.UserArgs] = None,
-    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+    env_vars: Mapping[str, Any] = immutabledict.immutabledict(),
 ) -> job_blocks.Packageable:
   # pyformat: disable
   """Dockerfile describes a Dockerfile for generating a docker image.