Write JSON state

Makefile

@@ -31,7 +31,7 @@ endif
 	@echo "make deb_bionic - Generate bionic deb package"
 	@echo "make itest_bionic - Run tests and integration checks"
 	@echo "make _itest_bionic - Run only integration checks"
-	@echo "make deb_jammy - Generate bionic deb package"
+	@echo "make deb_jammy - Generate jammy deb package"
 	@echo "make itest_jammy - Run tests and integration checks"
 	@echo "make _itest_jammy - Run only integration checks"
 	@echo "make release - Prepare debian info for new release"

bin/tronfig

@@ -199,7 +199,6 @@ if __name__ == "__main__":
     client = Client(args.server)
 
     if args.print_config:
-        # TODO: use maybe_encode()
         content = client.config(args.source)["config"]
         if type(content) is not bytes:
             content = content.encode("utf8")

docs/source/jobs.rst

@@ -209,7 +209,7 @@ Optional Fields
     after this duration.
 
 **trigger_downstreams** (bool or dict)
-    Upon successfull completion of an action, will emit a trigger for every
+    Upon successful completion of an action, will emit a trigger for every
     item in the dictionary. When set to ``true``, a default dict of
     ``{shortdate: "{shortdate}"}`` is assumed. Emitted triggers will be in form:
     ``<namespace>.<job>.<action>.<dict key>.<rendered value>``. See
@@ -220,7 +220,7 @@ Optional Fields
     have been emitted by upstream actions. Unlike with ``requires`` attribute,
     dependent actions don't have to belong to the same job. ``triggered_by``
     template may contain any pattern allowed in ``command`` attribute.
-    See :ref:`shortdate` for an explantion of shortdate
+    See :ref:`shortdate` for an explanation of shortdate
 
     Example:
 

itest.sh

@@ -60,14 +60,3 @@ fi
 
 kill -SIGTERM $TRON_PID
 wait $TRON_PID || true
-
-/opt/venvs/tron/bin/python - <<EOF
-import os
-from tron.serialize.runstate.shelvestore import ShelveStateStore, ShelveKey
-db = ShelveStateStore('$TRON_WORKDIR/tron_state')
-key = ShelveKey('mcp_state', 'StateMetadata')
-res = db.restore([key])
-ts = res[key][u'create_time']
-print("assert db time {} > start time {}".format(ts, int(os.environ['TRON_START_TIME'])))
-assert ts > int(os.environ['TRON_START_TIME'])
-EOF

requirements-dev-minimal.txt

@@ -8,6 +8,7 @@ pylint
 pytest
 pytest-asyncio
 requirements-tools
+types-pytz
 types-PyYAML
 types-requests<2.31.0.7 # newer types-requests requires urllib3>=2
 types-simplejson

requirements-dev.txt

@@ -29,6 +29,7 @@ requirements-tools==1.2.1
 toml==0.10.2
 tomli==2.0.1
 tomlkit==0.11.6
+types-pytz==2024.2.0.20240913
 types-PyYAML==6.0.12
 types-requests==2.31.0.5
 types-simplejson==3.19.0.20240310

tests/serialize/runstate/dynamodb_state_store_test.py

@@ -1,4 +1,4 @@
-import pickle
+import json
 from unittest import mock
 
 import boto3
@@ -100,24 +100,41 @@ def store():
 
 @pytest.fixture
 def small_object():
-    yield pickle.dumps({"this": "data"})
+    yield {
+        "job_name": "example_job",
+        "run_num": 1,
+        "run_time": None,
+        "node_name": "example_node",
+        "runs": [],
+        "cleanup_run": None,
+        "manual": False,
+    }
 
 
 @pytest.fixture
 def large_object():
-    yield pickle.dumps([i for i in range(1000000)])
+    yield {
+        "job_name": "example_job",
+        "run_num": 1,
+        "run_time": None,
+        "node_name": "example_node",
+        "runs": [],
+        "cleanup_run": None,
+        "manual": False,
+        "large_data": [i for i in range(1_000_000)],
+    }
 
 
 @pytest.mark.usefixtures("store", "small_object", "large_object")
 class TestDynamoDBStateStore:
     def test_save(self, store, small_object, large_object):
         key_value_pairs = [
             (
-                store.build_key("DynamoDBTest", "two"),
+                store.build_key("job_state", "two"),
                 small_object,
             ),
             (
-                store.build_key("DynamoDBTest2", "four"),
+                store.build_key("job_run_state", "four"),
                 small_object,
             ),
         ]
@@ -126,21 +143,27 @@ def test_save(self, store, small_object, large_object):
 
         assert store.save_errors == 0
         keys = [
-            store.build_key("DynamoDBTest", "two"),
-            store.build_key("DynamoDBTest2", "four"),
+            store.build_key("job_state", "two"),
+            store.build_key("job_run_state", "four"),
         ]
         vals = store.restore(keys)
         for key, value in key_value_pairs:
             assert_equal(vals[key], value)
 
+        for key in keys:
+            item = store.table.get_item(Key={"key": key, "index": 0})
+            assert "Item" in item
+            assert "json_val" in item["Item"]
+            assert_equal(json.loads(item["Item"]["json_val"]), small_object)
+
     def test_delete_if_val_is_none(self, store, small_object, large_object):
         key_value_pairs = [
             (
-                store.build_key("DynamoDBTest", "two"),
+                store.build_key("job_state", "two"),
                 small_object,
             ),
             (
-                store.build_key("DynamoDBTest2", "four"),
+                store.build_key("job_run_state", "four"),
                 small_object,
             ),
         ]
@@ -149,7 +172,7 @@ def test_delete_if_val_is_none(self, store, small_object, large_object):
 
         delete = [
             (
-                store.build_key("DynamoDBTest", "two"),
+                store.build_key("job_state", "two"),
                 None,
             ),
         ]
@@ -159,55 +182,55 @@ def test_delete_if_val_is_none(self, store, small_object, large_object):
         assert store.save_errors == 0
         # Try to restore both, we should just get one back
         keys = [
-            store.build_key("DynamoDBTest", "two"),
-            store.build_key("DynamoDBTest2", "four"),
+            store.build_key("job_state", "two"),
+            store.build_key("job_run_state", "four"),
         ]
         vals = store.restore(keys)
         assert vals == {keys[1]: small_object}
 
     def test_save_more_than_4KB(self, store, small_object, large_object):
         key_value_pairs = [
             (
-                store.build_key("DynamoDBTest", "two"),
+                store.build_key("job_state", "two"),
                 large_object,
             ),
         ]
         store.save(key_value_pairs)
         store._consume_save_queue()
 
         assert store.save_errors == 0
-        keys = [store.build_key("DynamoDBTest", "two")]
+        keys = [store.build_key("job_state", "two")]
         vals = store.restore(keys)
         for key, value in key_value_pairs:
             assert_equal(vals[key], value)
 
     def test_restore_more_than_4KB(self, store, small_object, large_object):
-        keys = [store.build_key("thing", i) for i in range(3)]
-        value = pickle.loads(large_object)
+        keys = [store.build_key("job_state", i) for i in range(3)]
+        value = large_object
         pairs = zip(keys, (value for i in range(len(keys))))
         store.save(pairs)
         store._consume_save_queue()
 
         assert store.save_errors == 0
         vals = store.restore(keys)
         for key in keys:
-            assert_equal(pickle.dumps(vals[key]), large_object)
+            assert_equal(vals[key], large_object)
 
     def test_restore(self, store, small_object, large_object):
-        keys = [store.build_key("thing", i) for i in range(3)]
-        value = pickle.loads(small_object)
+        keys = [store.build_key("job_state", i) for i in range(3)]
+        value = small_object
         pairs = zip(keys, (value for i in range(len(keys))))
         store.save(pairs)
         store._consume_save_queue()
 
         assert store.save_errors == 0
         vals = store.restore(keys)
         for key in keys:
-            assert_equal(pickle.dumps(vals[key]), small_object)
+            assert_equal(vals[key], small_object)
 
     def test_delete_item(self, store, small_object, large_object):
-        keys = [store.build_key("thing", i) for i in range(3)]
-        value = pickle.loads(large_object)
+        keys = [store.build_key("job_state", i) for i in range(3)]
+        value = large_object
         pairs = list(zip(keys, (value for i in range(len(keys)))))
         store.save(pairs)
 
@@ -222,8 +245,8 @@ def test_retry_saving(self, store, small_object, large_object):
             "moto.dynamodb2.responses.DynamoHandler.transact_write_items",
             side_effect=KeyError("foo"),
         ) as mock_failed_write:
-            keys = [store.build_key("thing", i) for i in range(1)]
-            value = pickle.loads(small_object)
+            keys = [store.build_key("job_state", i) for i in range(1)]
+            value = small_object
             pairs = zip(keys, (value for i in range(len(keys))))
             try:
                 store.save(pairs)
@@ -236,7 +259,7 @@ def test_retry_reading(self, store, small_object, large_object):
                 store.name: [
                     {
                         "index": {"N": "0"},
-                        "key": {"S": "thing 0"},
+                        "key": {"S": "job_state 0"},
                     },
                 ],
             },
@@ -246,15 +269,15 @@ def test_retry_reading(self, store, small_object, large_object):
                     "Keys": [
                         {
                             "index": {"N": "0"},
-                            "key": {"S": "thing 0"},
+                            "key": {"S": "job_state 0"},
                         }
                     ],
                 },
             },
             "ResponseMetadata": {},
         }
-        keys = [store.build_key("thing", i) for i in range(1)]
-        value = pickle.loads(small_object)
+        keys = [store.build_key("job_state", i) for i in range(1)]
+        value = small_object
         pairs = zip(keys, (value for i in range(len(keys))))
         store.save(pairs)
         with mock.patch.object(
@@ -269,7 +292,7 @@ def test_retry_reading(self, store, small_object, large_object):
 
     def test_restore_exception_propagation(self, store, small_object):
         # This test is to ensure that restore propagates exceptions upwards: see DAR-2328
-        keys = [store.build_key("thing", i) for i in range(3)]
+        keys = [store.build_key("job_state", i) for i in range(3)]
 
         mock_future = mock.MagicMock()
         mock_future.result.side_effect = Exception("mocked exception")

tests/serialize/runstate/statemanager_test.py

@@ -19,9 +19,7 @@
 from tron.serialize.runstate.statemanager import PersistenceStoreError
 from tron.serialize.runstate.statemanager import PersistentStateManager
 from tron.serialize.runstate.statemanager import StateChangeWatcher
-from tron.serialize.runstate.statemanager import StateMetadata
 from tron.serialize.runstate.statemanager import StateSaveBuffer
-from tron.serialize.runstate.statemanager import VersionMismatchError
 
 
 class TestPersistenceManagerFactory(TestCase):
@@ -42,24 +40,6 @@ def test_from_config_shelve(self):
             shutil.rmtree(tmpdir)
 
 
-class TestStateMetadata(TestCase):
-    def test_validate_metadata(self):
-        metadata = {"version": (0, 5, 2)}
-        StateMetadata.validate_metadata(metadata)
-
-    def test_validate_metadata_no_state_data(self):
-        metadata = None
-        StateMetadata.validate_metadata(metadata)
-
-    def test_validate_metadata_mismatch(self):
-        metadata = {"version": (200, 1, 1)}
-        assert_raises(
-            VersionMismatchError,
-            StateMetadata.validate_metadata,
-            metadata,
-        )
-
-
 class TestStateSaveBuffer(TestCase):
     @setup
     def setup_buffer(self):
@@ -103,15 +83,7 @@ def test_keys_for_items(self):
 
     def test_restore(self):
         job_names = ["one", "two"]
-        with mock.patch.object(
-            self.manager,
-            "_restore_metadata",
-            autospec=True,
-        ) as mock_restore_metadata, mock.patch.object(
-            self.manager,
-            "_restore_dicts",
-            autospec=True,
-        ) as mock_restore_dicts, mock.patch.object(
+        with mock.patch.object(self.manager, "_restore_dicts", autospec=True,) as mock_restore_dicts, mock.patch.object(
             self.manager,
             "_restore_runs_for_job",
             autospect=True,
@@ -125,7 +97,6 @@ def test_restore(self):
             ]
 
             restored_state = self.manager.restore(job_names)
-            mock_restore_metadata.assert_called_once_with()
             assert mock_restore_dicts.call_args_list == [
                 mock.call(runstate.JOB_STATE, job_names),
             ]
@@ -280,19 +251,6 @@ def test_save_job(self):
             mock_job.state_data,
         )
 
-    @mock.patch(
-        "tron.serialize.runstate.statemanager.StateMetadata",
-        autospec=None,
-    )
-    def test_save_metadata(self, mock_state_metadata):
-        self.watcher.save_metadata()
-        meta_data = mock_state_metadata.return_value
-        self.watcher.state_manager.save.assert_called_with(
-            runstate.MCP_STATE,
-            meta_data.name,
-            meta_data.state_data,
-        )
-
     def test_shutdown(self):
         self.watcher.shutdown()
         assert not self.watcher.state_manager.enabled