feat: Added no copy catalog (#229)

Author: vijayvammi

examples/04-catalog/catalog_no_copy.py

@@ -0,0 +1,50 @@
+"""
+You can execute this pipeline by:
+
+    python examples/04-catalog/catalog_python.py
+"""
+
+from examples.common.functions import check_files_do_not_exist, write_files
+from runnable import Catalog, Pipeline, PythonTask, ShellTask
+
+
+def main():
+    write_catalog = Catalog(put=["df.csv", "data_folder/data.txt"], store_copy=False)
+    generate_data = PythonTask(
+        name="generate_data_python",
+        function=write_files,
+        catalog=write_catalog,
+    )
+
+    delete_files_command = """
+        rm df.csv || true && \
+        rm data_folder/data.txt || true
+    """
+    # delete from local files after generate
+    # since its local catalog, we delete to show "get from catalog"
+    delete_local_after_generate = ShellTask(
+        name="delete_after_generate",
+        command=delete_files_command,
+    )
+
+    # Since store_copy was set to False, this step should fail
+    check_files_do_not_exist_task = PythonTask(
+        name="check_files_do_not_exist",
+        function=check_files_do_not_exist,
+        terminate_with_success=True,
+    )
+
+    pipeline = Pipeline(
+        steps=[
+            generate_data,
+            delete_local_after_generate,
+            check_files_do_not_exist_task,
+        ]
+    )
+    _ = pipeline.execute()
+
+    return pipeline
+
+
+if __name__ == "__main__":
+    main()

examples/11-jobs/catalog_no_copy.py

@@ -0,0 +1,18 @@
+from examples.common.functions import write_files
+from runnable import Catalog, PythonJob
+
+
+def main():
+    write_catalog = Catalog(put=["df.csv", "data_folder/data.txt"], store_copy=False)
+    job = PythonJob(
+        function=write_files,
+        catalog=write_catalog,
+    )
+
+    job.execute()
+
+    return job
+
+
+if __name__ == "__main__":
+    main()

examples/common/functions.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from pathlib import Path
 from typing import Dict, List, Union
 
@@ -124,6 +125,12 @@ def read_files():
     assert data.strip() == "hello world"
 
 
+def check_files_do_not_exist():
+    run_id = os.environ.get("RUNNABLE_RUN_ID")
+    assert not Path(f".catalog/{run_id}/df.csv").exists()
+    assert not Path(f".catalog/{run_id}/data_folder/data.txt").exists()
+
+
 def process_chunk(chunk: int):
     """
     An example function that processes a chunk of data.

extensions/catalog/any_path.py

@@ -95,7 +95,10 @@ def get(self, name: str) -> List[DataCatalog]:
         return data_catalogs
 
     def put(
-        self, name: str, allow_file_not_found_exc: bool = False
+        self,
+        name: str,
+        allow_file_not_found_exc: bool = False,
+        store_copy: bool = True,
     ) -> List[DataCatalog]:
         """
         Put the files matching the glob pattern into the catalog.
@@ -154,7 +157,15 @@ def put(
             data_catalogs.append(data_catalog)
 
             # TODO: Think about syncing only if the file is changed
-            self.upload_to_catalog(file)
+            if store_copy:
+                logger.debug(
+                    f"Copying file {file} to the catalog location for run_id: {run_id}"
+                )
+                self.upload_to_catalog(file)
+            else:
+                logger.debug(
+                    f"Not copying file {file} to the catalog location for run_id: {run_id}"
+                )
 
         if not data_catalogs and not allow_file_not_found_exc:
             raise Exception(f"Did not find any files matching {name} in {copy_from}")

extensions/job_executor/__init__.py

@@ -29,6 +29,7 @@ class GenericJobExecutor(BaseJobExecutor):
     @property
     def _context(self):
         assert context.run_context
+        assert isinstance(context.run_context, context.JobContext)
         return context.run_context
 
     def _get_parameters(self) -> Dict[str, JsonParameter]:
@@ -147,7 +148,9 @@ def _sync_catalog(
         data_catalogs = []
         for name_pattern in catalog_settings:
             data_catalog = self._context.catalog.put(
-                name=name_pattern, allow_file_not_found_exc=allow_file_not_found_exc
+                name=name_pattern,
+                allow_file_not_found_exc=allow_file_not_found_exc,
+                store_copy=self._context.catalog_store_copy,
             )
 
             logger.debug(f"Added data catalog: {data_catalog} to job log")

extensions/pipeline_executor/__init__.py

@@ -160,7 +160,9 @@ def _sync_catalog(
 
             elif stage == "put":
                 data_catalog = self._context.catalog.put(
-                    name=name_pattern, allow_file_not_found_exc=allow_file_no_found_exc
+                    name=name_pattern,
+                    allow_file_not_found_exc=allow_file_no_found_exc,
+                    store_copy=node_catalog_settings.get("store_copy", True),
                 )
             else:
                 raise Exception(f"Stage {stage} not supported")

runnable/catalog.py

@@ -57,7 +57,7 @@ def get(self, name: str) -> List[DataCatalog]:
 
     @abstractmethod
     def put(
-        self, name: str, allow_file_not_found_exc: bool = False
+        self, name: str, allow_file_not_found_exc: bool = False, store_copy: bool = True
     ) -> List[DataCatalog]:
         """
         Put the file by 'name' from the 'compute_data_folder' in the catalog for the run_id.
@@ -120,7 +120,10 @@ def get(self, name: str) -> List[DataCatalog]:
         return []
 
     def put(
-        self, name: str, allow_file_not_found_exc: bool = False
+        self,
+        name: str,
+        allow_file_not_found_exc: bool = False,
+        store_copy: bool = True,
     ) -> List[DataCatalog]:
         """
         Does nothing

runnable/context.py

@@ -475,6 +475,7 @@ class JobContext(RunnableContext):
         default=None,
         description="Catalog settings to be used for the job.",
     )
+    catalog_store_copy: bool = Field(default=True, alias="catalog_store_copy")
 
     @computed_field  # type: ignore
     @cached_property

runnable/nodes.py

@@ -411,11 +411,13 @@ def _get_executor_config(self, executor_type) -> str:
         return self.overrides.get(executor_type) or ""
 
 
+# Unfortunately, this is defined in 2 places. Look in SDK
 class CatalogStructure(BaseModel):
     model_config = ConfigDict(extra="forbid")  # Need to forbid
 
     get: List[str] = Field(default_factory=list)
     put: List[str] = Field(default_factory=list)
+    store_copy: bool = Field(default=True, alias="store_copy")
 
 
 class ExecutableNode(TraversalNode):

runnable/sdk.py

@@ -60,6 +60,7 @@ class Catalog(BaseModel):
     Attributes:
         get (List[str]): List of glob patterns to get from central catalog to the compute data folder.
         put (List[str]): List of glob patterns to put into central catalog from the compute data folder.
+        store_copy (bool): Whether to store a copy of the data in the central catalog.
 
     Examples:
         >>> from runnable import Catalog
@@ -74,6 +75,7 @@ class Catalog(BaseModel):
     # compute_data_folder: str = Field(default="", alias="compute_data_folder")
     get: List[str] = Field(default_factory=list, alias="get")
     put: List[str] = Field(default_factory=list, alias="put")
+    store_copy: bool = Field(default=True, alias="store_copy")
 
 
 class BaseTraversal(ABC, BaseModel):
@@ -845,6 +847,11 @@ def return_catalog_settings(self) -> Optional[List[str]]:
             return []
         return self.catalog.put
 
+    def return_bool_catalog_store_copy(self) -> bool:
+        if self.catalog is None:
+            return True
+        return self.catalog.store_copy
+
     def _is_called_for_definition(self) -> bool:
         """
         If the run context is set, we are coming in only to get the pipeline definition.
@@ -888,6 +895,7 @@ def execute(
         }
 
         run_context = context.JobContext.model_validate(configurations)
+        run_context.catalog_store_copy = self.return_bool_catalog_store_copy()
 
         assert isinstance(run_context.job_executor, BaseJobExecutor)