Use pickle instead of serialize (#29)

* Use pickle instead of serialize

* Fix path to pkl

* Add func to generate pickle path

* Fmt

* Fix target

* Fix target

* Remove unused import

* Change pkl path

* Update example

* fmt

---------

Co-authored-by: yutaro-oguri <[email protected]>

Author: maronuu

example/run_child.py

@@ -4,18 +4,17 @@
 import fire
 import gokart
 import luigi
-# Import task definition
-import tasks  # noqa: F401
+from gokart.target import make_target
 
 logging.basicConfig(level=logging.INFO)
 
 
-def main(serialized_task: str) -> None:
+def main(task_pkl_path: str) -> None:
     # Load luigi config
     luigi.configuration.LuigiConfigParser.add_config_path("./conf/base.ini")
 
     # Parse a serialized gokart.TaskOnKart
-    task: gokart.TaskOnKart = gokart.TaskInstanceParameter().parse(serialized_task)
+    task: gokart.TaskOnKart = make_target(task_pkl_path).load()
 
     # Run gokart.build
     gokart.build(task)

kannon/master.py

@@ -6,6 +6,7 @@
 from typing import Deque, Dict, List, Optional, Set
 
 import gokart
+from gokart.target import make_target
 from kubernetes import client
 from luigi.task import flatten
 
@@ -107,25 +108,27 @@ def _exec_gokart_task(self, task: gokart.TaskOnKart) -> None:
             raise RuntimeError(f"Task {self._gen_task_info(task)} on job master has failed.")
 
     def _exec_bullet_task(self, task: TaskOnBullet) -> None:
+        # Save task instance as pickle object
+        pkl_path = self._gen_pkl_path(task)
+        make_target(pkl_path).dump(task)
         # Run on child job
-        serialized_task = gokart.TaskInstanceParameter().serialize(task)
         job_name = gen_job_name(f"{self.job_prefix}-{task.get_task_family()}")
         job = self._create_child_job_object(
             job_name=job_name,
-            serialized_task=serialized_task,
+            task_pkl_path=pkl_path,
         )
         create_job(self.api_instance, job, self.namespace)
         logger.info(f"Created child job {job_name} with task {self._gen_task_info(task)}")
         task_unique_id = task.make_unique_id()
         self.task_id_to_job_name[task_unique_id] = job_name
 
-    def _create_child_job_object(self, job_name: str, serialized_task: str) -> client.V1Job:
+    def _create_child_job_object(self, job_name: str, task_pkl_path: str) -> client.V1Job:
         # TODO: use python -c to avoid dependency to execute_task.py
         cmd = [
             "python",
             self.path_child_script,
-            "--serialized-task",
-            f"'{serialized_task}'",
+            "--task-pkl-path",
+            f"'{task_pkl_path}'",
         ]
         job = deepcopy(self.template_job)
         # replace command
@@ -150,6 +153,10 @@ def _create_child_job_object(self, job_name: str, serialized_task: str) -> clien
     def _gen_task_info(task: gokart.TaskOnKart) -> str:
         return f"{task.get_task_family()}_{task.make_unique_id()}"
 
+    @staticmethod
+    def _gen_pkl_path(task: gokart.TaskOnKart) -> str:
+        return os.path.join(task.workspace_directory, 'kannon', f'task_obj_{task.make_unique_id()}.pkl')
+
     def _is_executable(self, task: gokart.TaskOnKart) -> bool:
         children = flatten(task.requires())
 

test/test_master.py

@@ -76,7 +76,7 @@ def test_success_basic(self) -> None:
         class Example(gokart.TaskOnKart):
             pass
 
-        serialized_task = gokart.TaskInstanceParameter().serialize(Example())
+        path_to_pkl = "path/to/obj"
         template_job = self._get_template_job()
         master = Kannon(
             api_instance=None,
@@ -88,7 +88,7 @@ class Example(gokart.TaskOnKart):
         # set os env
         os.environ.update({"TASK_WORKSPACE_DIRECTORY": "/cache"})
         child_job_name = "test-job"
-        child_job = master._create_child_job_object(child_job_name, serialized_task)
+        child_job = master._create_child_job_object(child_job_name, path_to_pkl)
 
         # following should be copied from template_job
         self.assertEqual(child_job.api_version, template_job.api_version)
@@ -99,7 +99,7 @@ class Example(gokart.TaskOnKart):
         self.assertEqual(child_job.spec.template.spec.containers[0].image, template_job.spec.template.spec.containers[0].image)
         self.assertEqual(child_job.spec.template.spec.restart_policy, template_job.spec.template.spec.restart_policy)
         # following should be overwritten
-        self.assertEqual(child_job.spec.template.spec.containers[0].command, ["python", __file__, "--serialized-task", f"'{serialized_task}'"])
+        self.assertEqual(child_job.spec.template.spec.containers[0].command, ["python", __file__, "--task-pkl-path", f"'{path_to_pkl}'"])
         self.assertEqual(child_job.metadata.name, child_job_name)
         # envvar TASK_WORKSPACE_DIRECTORY should be inherited
         child_env = child_job.spec.template.spec.containers[0].env
@@ -111,7 +111,7 @@ def test_success_custom_env(self) -> None:
         class Example(gokart.TaskOnKart):
             pass
 
-        serialized_task = gokart.TaskInstanceParameter().serialize(Example())
+        path_to_pkl = "path/to/obj"
         template_job = self._get_template_job()
         master = Kannon(
             api_instance=None,
@@ -123,7 +123,7 @@ class Example(gokart.TaskOnKart):
         # set os env
         os.environ.update({"TASK_WORKSPACE_DIRECTORY": "/cache", "MY_ENV0": "env0", "MY_ENV1": "env1"})
         child_job_name = "test-job"
-        child_job = master._create_child_job_object(child_job_name, serialized_task)
+        child_job = master._create_child_job_object(child_job_name, path_to_pkl)
 
         child_env = child_job.spec.template.spec.containers[0].env
         self.assertEqual(len(child_env), 3)
@@ -136,7 +136,7 @@ def test_fail_command_set(self) -> None:
         class Example(gokart.TaskOnKart):
             pass
 
-        serialized_task = gokart.TaskInstanceParameter().serialize(Example())
+        path_to_pkl = "path/to/obj"
         template_job = self._get_template_job()
         template_job.spec.template.spec.containers[0].command = ["dummy-command"]
         master = Kannon(
@@ -149,14 +149,14 @@ class Example(gokart.TaskOnKart):
         # set os env
         os.environ.update({"TASK_WORKSPACE_DIRECTORY": "/cache"})
         with self.assertRaises(AssertionError):
-            master._create_child_job_object("test-job", serialized_task)
+            master._create_child_job_object("test-job", path_to_pkl)
 
     def test_fail_default_env_not_exist(self) -> None:
 
         class Example(gokart.TaskOnKart):
             pass
 
-        serialized_task = gokart.TaskInstanceParameter().serialize(Example())
+        path_to_pkl = "path/to/obj"
         template_job = self._get_template_job()
 
         cases = [None, ["TASK_WORKSPACE_DIRECTORY", "MY_ENV0", "MY_ENV1"]]
@@ -170,7 +170,7 @@ class Example(gokart.TaskOnKart):
                     env_to_inherit=case,
                 )
                 with self.assertRaises(ValueError):
-                    master._create_child_job_object("test-job", serialized_task)
+                    master._create_child_job_object("test-job", path_to_pkl)
 
 
 if __name__ == '__main__':