RWM is unable to process larger input files split into chunks (#53)

* RWM is unable to process larger input files split into chunks

Author: AndrewNepogoda

cli/README.md

@@ -6,7 +6,7 @@ This Command-Line Interface (CLI) is designed to provide an easy and interactive
 1. Create a batch configuration (ex. `poc.json`) file using the syntax and structure outlined in the [RAI Workflow Framework README](../workflow/README.md).
 2. Add `rai-workflow-manager` as dependency to your `requirements.txt` file:
 ```txt
-rai-workflow-manager==0.0.28
+rai-workflow-manager==0.0.29
 ```
 3. Build the project:
 ```bash

workflow/__init__.py

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version_info__ = (0, 0, 28)
+__version_info__ = (0, 0, 29)
 __version__ = ".".join(map(str, __version_info__))

workflow/executor.py

@@ -351,7 +351,7 @@ def _execute(self, logger: logging.Logger, env_config: EnvConfig, rai_config: Ra
     def await_pending(self, env_config, logger, missed_resources):
         loop = asyncio.get_event_loop()
         if loop.is_running():
-            raise Exception('Waiting for resource would interrupt unexpected event loop - aborting to avoid confusion.')
+            raise Exception('Waiting for resource would interrupt unexpected event loop - aborting to avoid confusion')
         pending = [src for src in missed_resources if self._resource_is_async(src)]
         pending_cos = [self._await_async_resource(logger, env_config, resource) for resource in pending]
         loop.run_until_complete(asyncio.gather(*pending_cos))
@@ -365,28 +365,35 @@ async def _await_async_resource(self, logger: logging.Logger, env_config: EnvCon
     def _load_source(self, logger: logging.Logger, env_config: EnvConfig, rai_config: RaiConfig, src):
         source_name = src["source"]
         if 'is_date_partitioned' in src and src['is_date_partitioned'] == 'Y':
+            logger.info(f"Loading source '{source_name}' partitioned by date")
             if self.collapse_partitions_on_load:
                 srcs = src["dates"]
                 first_date = srcs[0]["date"]
                 last_date = srcs[-1]["date"]
 
                 logger.info(
-                    f"Loading '{source_name}' from all partitions simultaneously, range {first_date} to {last_date}")
+                    f"Loading '{source_name}' all date partitions simultaneously, range {first_date} to {last_date}")
 
                 resources = []
                 for d in srcs:
                     resources += d["resources"]
                 self._load_resource(logger, env_config, rai_config, resources, src)
             else:
-                logger.info(f"Loading '{source_name}' one partition at a time")
+                logger.info(f"Loading '{source_name}' one date partition at a time")
                 for d in src["dates"]:
                     logger.info(f"Loading partition for date {d['date']}")
 
-                    resources = d["resources"]
-                    self._load_resource(logger, env_config, rai_config, resources, src)
+                    for res in d["resources"]:
+                        self._load_resource(logger, env_config, rai_config, [res], src)
         else:
-            logger.info(f"Loading source '{source_name}' not partitioned by date ")
-            self._load_resource(logger, env_config, rai_config, src["resources"], src)
+            logger.info(f"Loading source '{source_name}' not partitioned by date")
+            if self.collapse_partitions_on_load:
+                logger.info(f"Loading '{source_name}' all chunk partitions simultaneously")
+                self._load_resource(logger, env_config, rai_config, src["resources"], src)
+            else:
+                logger.info(f"Loading '{source_name}' one chunk partition at a time")
+                for res in src["resources"]:
+                    self._load_resource(logger, env_config, rai_config, [res], src)
 
     @staticmethod
     def _resource_is_async(src):
@@ -539,7 +546,7 @@ def _load_exports(logger: logging.Logger, env_config: EnvConfig, src) -> List[Ex
                                           container=env_config.get_container(e.get("container", default_container)),
                                           offset_by_number_of_days=e.get("offsetByNumberOfDays", 0)))
                 except KeyError as ex:
-                    logger.warning(f"Unsupported FileType: {ex}. Skipping export: {e}.")
+                    logger.warning(f"Unsupported FileType: {ex}. Skipping export: {e}")
         return exports
 
 

workflow/manager.py

@@ -87,7 +87,7 @@ def add_engine(self, size: str = "XS") -> None:
         :param size:    RAI engine size
         :return:
         """
-        self.__logger.info(f"Trying to add engine with `{size}` size to manager.")
+        self.__logger.info(f"Trying to add engine with `{size}` size to manager")
         config = self.get_rai_config(size)
         if self.__engines:
             if size in self.__engines:
@@ -108,7 +108,7 @@ def remove_engine(self, size: str = "XS") -> None:
         :param size:    RAI engine size
         :return:
         """
-        self.__logger.info(f"Trying to remove engine with `{size}` size from manager.")
+        self.__logger.info(f"Trying to remove engine with `{size}` size from manager")
         if size not in self.__engines:
             self.__logger.info(f"`{size}` isn't managed. Ignore deletion")
         else:
@@ -128,7 +128,7 @@ def provision_engine(self, size: str) -> None:
         :param size:    RAI engine size
         :return:
         """
-        self.__logger.info(f"Trying to provision engine with `{size}` size and add to manager.")
+        self.__logger.info(f"Trying to provision engine with `{size}` size and add to manager")
         config = self.get_rai_config(size)
         if self.__engines:
             if size in self.__engines:
@@ -141,7 +141,7 @@ def provision_engine(self, size: str) -> None:
                 self.__create_engine(config, size)
                 self.__engines[size] = EngineMetaInfo(config.engine, size, False)
         else:
-            self.__logger.info(f"Provision engine `{config.engine}` as default.")
+            self.__logger.info(f"Provision engine `{config.engine}` as default")
             self.__recreate_engine(config, size)
             self.__engines = {size: EngineMetaInfo(config.engine, size, True)}
 

workflow/query.py

@@ -151,10 +151,10 @@ def load_resources(logger: logging.Logger, config: AzureConfig, resources, src)
             logger.error(f"Unknown file type {file_stype_str}")
     else:
         if src_type == ContainerType.LOCAL:
-            logger.info("Loading from local file.")
+            logger.info("Loading from local file")
             return _local_load_simple_query(rel_name, resources[0]["uri"], file_type)
         elif src_type == ContainerType.AZURE:
-            logger.info("Loading from Azure file.")
+            logger.info("Loading from Azure file")
             return _azure_load_simple_query(rel_name, resources[0]["uri"], file_type, config)