Add support for log files.

PiperOrigin-RevId: 378952727
Change-Id: I99aeb24a817582f36c66b866594a9de084c99ff3

Author: Stewart Miles

service/api/falken_service.py

@@ -233,6 +233,7 @@ def main(argv):
   if len(argv) > 1:
     logging.error('Non-flag parameters are not allowed: %s', argv)
   logging.debug('Starting the API service...')
+  logging.get_absl_handler().use_absl_log_file()
   server = serve()
   server.wait_for_termination()
 

service/launcher.py

@@ -102,7 +102,8 @@ def run_api(current_path: str):
   args = [
       sys.executable, '-m', 'api.falken_service', '--root_dir', FLAGS.root_dir,
       '--port', str(FLAGS.port), '--ssl_dir', FLAGS.ssl_dir,
-      '--verbosity', str(FLAGS.verbosity), '--alsologtostderr'
+      '--verbosity', str(FLAGS.verbosity), '--alsologtostderr',
+      '--log_dir', FLAGS.log_dir,
   ]
   for hyperparameters in FLAGS.hyperparameters:
     args.extend(['--hyperparameters', hyperparameters])
@@ -125,7 +126,7 @@ def run_learner(current_path: str):
   return subprocess.Popen(
       [sys.executable, '-m', 'learner.learner_service',
        '--root_dir', FLAGS.root_dir, '--verbosity', str(FLAGS.verbosity),
-       '--alsologtostderr'],
+       '--alsologtostderr', '--log_dir', FLAGS.log_dir],
       env=os.environ, cwd=current_path)
 
 
@@ -144,6 +145,7 @@ def main(argv):
   if len(argv) > 1:
     logging.error('Non-flag parameters are not allowed.')
 
+  logging.get_absl_handler().use_absl_log_file()
   check_ssl()
 
   file_dir = os.path.dirname(os.path.abspath(__file__))

service/launcher_test.py

@@ -21,6 +21,7 @@
 import tempfile
 from unittest import mock
 
+from absl import logging
 from absl.testing import absltest
 import common.generate_protos  # pylint: disable=unused-import
 import launcher
@@ -37,6 +38,7 @@ def setUp(self):
     os.makedirs(launcher.FLAGS.ssl_dir)
     launcher.FLAGS.root_dir = os.path.join(self.temp_dir.name, 'root_dir')
     os.makedirs(launcher.FLAGS.root_dir)
+    logging.FLAGS.log_dir = self.temp_dir
 
   def tearDown(self):
     """Tear down the testing environment."""
@@ -55,6 +57,7 @@ def test_run_api_test(self, popen):
          '--root_dir', launcher.FLAGS.root_dir, '--port', '50051',
          '--ssl_dir', launcher.FLAGS.ssl_dir,
          '--verbosity', '0', '--alsologtostderr',
+         '--log_dir', self.temp_dir,
          '--hyperparameters', '{"foo": 1}',
          '--hyperparameters', '{"bar": 1000}'],
         env=os.environ, cwd='mock_path')
@@ -81,7 +84,7 @@ def test_run_learner_test(self, popen):
     popen.assert_called_once_with(
         [sys.executable, '-m', 'learner.learner_service',
          '--root_dir', launcher.FLAGS.root_dir, '--verbosity', '0',
-         '--alsologtostderr'],
+         '--alsologtostderr', '--log_dir', self.temp_dir],
         env=os.environ, cwd='mock_path')
 
 

service/learner/learner_service.py

@@ -23,6 +23,7 @@
 
 from absl import app
 from absl import flags
+from absl import logging
 from data_store import data_store as data_store_module
 from data_store import file_system as data_store_file_system
 from learner import learner as learner_module
@@ -137,6 +138,7 @@ def shutdown(self):
 
 def main(unused_argv, assignment_path=None, iterations=-1):
   """Start a learner from the module's flags."""
+  logging.get_absl_handler().use_absl_log_file()
   service = LearnerService(assignment_path)
   # Gracefully shut down when handling SIGTERM.
   signal.signal(signal.SIGTERM,