remote/keras_remote/core/core_test.py at 59475d77c78ae99c29903e9b23836b41c736e9d9 · keras-team/remote · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
"""Tests for keras_remote.core.core — run decorator and env var capture."""

import os
from unittest import mock
from unittest.mock import MagicMock

from absl.testing import absltest

from keras_remote.core.core import run


class TestRunDecorator(absltest.TestCase):
  def test_decorator_preserves_wrapped_function(self):
    @run(accelerator="cpu")
    def my_func():
      """My docstring."""

    self.assertTrue(callable(my_func))
    self.assertEqual(my_func.__name__, "my_func")
    self.assertEqual(my_func.__doc__, "My docstring.")


class TestEnvVarCapture(absltest.TestCase):
  def test_exact_match(self):
    with (
      mock.patch.dict(os.environ, {"MY_VAR": "my_val"}),
      mock.patch("keras_remote.core.core._execute_on_gke") as mock_exec,
    ):

      @run(accelerator="cpu", capture_env_vars=["MY_VAR"])
      def func():
        pass

      func()
      call_args = mock_exec.call_args
      env_vars = call_args[0][-1]  # last positional arg
      self.assertEqual(env_vars, {"MY_VAR": "my_val"})

  def test_wildcard_pattern(self):
    env = {
      k: v
      for k, v in os.environ.items()
      if k not in ("PREFIX_A", "PREFIX_B", "OTHER")
    }
    env.update({"PREFIX_A": "1", "PREFIX_B": "2", "OTHER": "3"})
    with (
      mock.patch.dict(os.environ, env, clear=True),
      mock.patch("keras_remote.core.core._execute_on_gke") as mock_exec,
    ):

      @run(accelerator="cpu", capture_env_vars=["PREFIX_*"])
      def func():
        pass

      func()
      env_vars = mock_exec.call_args[0][-1]
      self.assertIn("PREFIX_A", env_vars)
      self.assertIn("PREFIX_B", env_vars)
      self.assertNotIn("OTHER", env_vars)

  def test_missing_var_skipped(self):
    env = {k: v for k, v in os.environ.items() if k != "NONEXISTENT"}
    with (
      mock.patch.dict(os.environ, env, clear=True),
      mock.patch("keras_remote.core.core._execute_on_gke") as mock_exec,
    ):

      @run(accelerator="cpu", capture_env_vars=["NONEXISTENT"])
      def func():
        pass

      func()
      env_vars = mock_exec.call_args[0][-1]
      self.assertEqual(env_vars, {})

  def test_none_capture(self):
    with mock.patch("keras_remote.core.core._execute_on_gke") as mock_exec:

      @run(accelerator="cpu", capture_env_vars=None)
      def func():
        pass

      func()
      env_vars = mock_exec.call_args[0][-1]
      self.assertEqual(env_vars, {})

  def test_mixed_exact_and_wildcard(self):
    env = {
      k: v
      for k, v in os.environ.items()
      if k not in ("EXACT_VAR", "WILD_A", "WILD_B")
    }
    env.update({"EXACT_VAR": "exact", "WILD_A": "a", "WILD_B": "b"})
    with (
      mock.patch.dict(os.environ, env, clear=True),
      mock.patch("keras_remote.core.core._execute_on_gke") as mock_exec,
    ):

      @run(
        accelerator="cpu",
        capture_env_vars=["EXACT_VAR", "WILD_*"],
      )
      def func():
        pass

      func()
      env_vars = mock_exec.call_args[0][-1]
      self.assertEqual(
        env_vars, {"EXACT_VAR": "exact", "WILD_A": "a", "WILD_B": "b"}
      )


class TestExecuteOnGkeDefaults(absltest.TestCase):
  def test_cluster_from_env(self):
    """When cluster=None, falls back to KERAS_REMOTE_CLUSTER env var."""
    with (
      mock.patch.dict(
        os.environ,
        {
          "KERAS_REMOTE_CLUSTER": "env-cluster",
          "KERAS_REMOTE_PROJECT": "proj",
        },
      ),
      mock.patch(
        "keras_remote.core.core.execute_remote",
        return_value=42,
      ) as mock_exec,
      mock.patch(
        "keras_remote.core.core.JobContext.from_params",
        return_value=MagicMock(),
      ),
    ):

      @run(accelerator="cpu", cluster=None)
      def func():
        pass

      func()

      call_args = mock_exec.call_args
      backend = call_args[0][1]
      self.assertEqual(backend.cluster, "env-cluster")

  def test_namespace_from_env(self):
    """When namespace=None, falls back to KERAS_REMOTE_GKE_NAMESPACE env var."""
    with (
      mock.patch.dict(
        os.environ,
        {
          "KERAS_REMOTE_GKE_NAMESPACE": "custom-ns",
          "KERAS_REMOTE_PROJECT": "proj",
        },
      ),
      mock.patch(
        "keras_remote.core.core.execute_remote",
        return_value=42,
      ) as mock_exec,
      mock.patch(
        "keras_remote.core.core.JobContext.from_params",
        return_value=MagicMock(),
      ),
    ):

      @run(accelerator="cpu", namespace=None)
      def func():
        pass

      func()

      backend = mock_exec.call_args[0][1]
      self.assertEqual(backend.namespace, "custom-ns")


if __name__ == "__main__":
  absltest.main()