Packaging tools jar with the python package

user_tools/build.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2023-2024, NVIDIA CORPORATION.
+# Copyright (c) 2023-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,29 +14,74 @@
 # limitations under the License.
 
 
-# Usage: ./build.sh [build_mode]
-# This script takes an optional build_mode ("fat" or otherwise) as a parameter.
-# If the build_mode is "fat", it runs the dependency downloader script and build the jar from
-# source.
+# Usage: ./build.sh [build_mode] [jar_url]
+# This script takes a build_mode ("fat" or "non-fat") as a parameter.
+# If the build_mode is "fat", it runs the dependency downloader script and build the jar from source.
 # If the build_mode is not "fat", it removes dependencies from the prepackaged directory
+# If jar_url is provided, the script downloads the JAR from the URL instead of building it
 # Finally performs the default build process.
 
 # Get the build mode argument
 build_mode="$1"
+jar_url="$2"
 
 # get the directory of the script
 WORK_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]:-$0}"; )" &> /dev/null && pwd 2> /dev/null; )";
 
 # Define resource directory
 RESOURCE_DIR="src/spark_rapids_pytools/resources"
+TOOLS_RESOURCE_FOLDER="tools-resources"
 PREPACKAGED_FOLDER="csp-resources"
 
 # Constants and variables of core module
 CORE_DIR="$WORK_DIR/../core"
 TOOLS_JAR_FILE=""
+DOWNLOAD_DIR="$WORK_DIR/downloaded_jars"
 
+# Function to download JAR from URL
+download_jar_from_url() {
+  local url="$1"
+  local output_dir="$2"
+
+  # Create download directory if it doesn't exist
+  mkdir -p "$output_dir"
+
+  # Extract filename from URL
+  local filename=$(basename "$url")
+  local output_path="$output_dir/$filename"
+
+  # Validate the JAR filename to ensure it is not a source, javadoc, or test JAR
+  if [[ "$filename" == *"sources.jar" || "$filename" == *"javadoc.jar" || "$filename" == *"tests.jar" ]]; then
+    echo "Invalid JAR file: $filename. Source, javadoc, or test JARs are not allowed."
+    exit 1
+  fi
+
+  echo "Downloading JAR from $url to $output_path"
+  curl -L -o "$output_path" "$url"
+
+  if [ $? -ne 0 ]; then
+    echo "Failed to download JAR from URL: $url"
+    exit 1
+  fi
+
+  TOOLS_JAR_FILE="$output_path"
+  CLEANUP_JAR_FILE="$output_path"
+  echo "Downloaded JAR file: $TOOLS_JAR_FILE"
+}
+
+clean_up_downloaded_jar() {
+  # Check if the file exists
+  if [ -f "$CLEANUP_JAR_FILE" ]; then
+    echo "Cleaning up downloaded JAR file: $CLEANUP_JAR_FILE"
+    rm -f "$CLEANUP_JAR_FILE"
+  else
+    echo "No downloaded JAR file to clean up."
+  fi
+}
 
 # Function to run mvn command to build the tools jar
+# This function skips the test cases and builds the jar file and only
+# picks the jar file without sources/javadoc/tests..
 build_jar_from_source() {
   # store teh current directory
   local curr_dir
@@ -62,11 +107,22 @@ build_jar_from_source() {
   cd "$curr_dir" || exit
 }
 
-# Function to run the dependency downloader script for fat mode
+# Function to run the dependency downloader script for non-fat/fat mode
+# prepackage_mgr.py file downloads the dependencies for the csp-related resources
+# in case of fat mode.
+# In case of non-fat mode, it just copies the tools jar into the tools-resources folder
+# --fetch_all_csp=True toggles the fat/non-fat mode for the script
 download_web_dependencies() {
   local res_dir="$1"
+  local is_fat_mode="$2"
   local web_downloader_script="$res_dir/dev/prepackage_mgr.py"
-  python "$web_downloader_script" run --resource_dir="$res_dir" --tools_jar="$TOOLS_JAR_FILE"
+  if [ "$is_fat_mode" = "true" ]; then
+    echo "Downloading dependencies for fat mode"
+    python "$web_downloader_script" run --resource_dir="$res_dir" --tools_jar="$TOOLS_JAR_FILE" --fetch_all_csp=True
+  else
+    echo "Downloading dependencies for non-fat mode"
+    python "$web_downloader_script" run --resource_dir="$res_dir" --tools_jar="$TOOLS_JAR_FILE" --fetch_all_csp=False
+  fi
   if [ $? -ne 0 ]; then
     echo "Dependency download failed for fat mode. Exiting"
     exit 1
@@ -76,6 +132,8 @@ download_web_dependencies() {
 # Function to remove dependencies from the fat directory
 remove_web_dependencies() {
   local res_dir="$1"
+  # remove tools jar
+  rm -rf "${res_dir:?}"/"$TOOLS_RESOURCE_FOLDER"
   # remove folder recursively
   rm -rf "${res_dir:?}"/"$PREPACKAGED_FOLDER"
   # remove compressed file in case archive-mode was enabled
@@ -90,18 +148,36 @@ pre_build() {
 
 # Build process
 build() {
+  # Deletes pre-existing csp-resources.tgz folder
   remove_web_dependencies "$RESOURCE_DIR"
+  # Build the tools jar from source
+  if [ -n "$jar_url" ]; then
+      echo "Using provided JAR URL instead of building from source"
+      download_jar_from_url "$jar_url" "$DOWNLOAD_DIR"
+    else
+      echo "Building JAR from source"
+      build_jar_from_source
+  fi
   if [ "$build_mode" = "fat" ]; then
     echo "Building in fat mode"
-    build_jar_from_source
-    download_web_dependencies "$RESOURCE_DIR"
+    # This will download the dependencies and create the csp-resources
+    # and copy the dependencies into the csp-resources folder
+    # Tools resources are copied into the tools-resources folder
+    download_web_dependencies "$RESOURCE_DIR" "true"
+  else
+    echo "Building in non-fat mode"
+    # This will just copy the tools jar built from source into the tools-resources folder
+    download_web_dependencies "$RESOURCE_DIR" "false"
   fi
+  # Builds the python wheel file
+  # Look into the pyproject.toml file for the build system requirements
   python -m build --wheel
+  clean_up_downloaded_jar
 }
 
 # Main script execution
 pre_build
-build "$build_mode"
+build
 
 # Check build status
 if [ $? -eq 0 ]; then

user_tools/pyproject.toml

@@ -102,7 +102,7 @@ where = ["src"]
 [tool.setuptools.dynamic]
 version = {attr = "spark_rapids_pytools.__version__"}
 [tool.setuptools.package-data]
-"*"= ["*.json", "*.yaml", "*.ms", "*.sh", "*.tgz", "*.properties"]
+"*"= ["*.json", "*.yaml", "*.ms", "*.sh", "*.tgz", "*.properties","*.jar"]
 [tool.poetry]
 repository = "https://github.com/NVIDIA/spark-rapids-tools/tree/main"
 [project.optional-dependencies]

user_tools/src/spark_rapids_pytools/rapids/rapids_tool.py

@@ -427,7 +427,6 @@ class RapidsJarTool(RapidsTool):
 
     def _process_jar_arg(self):
         # TODO: use the StorageLib to download the jar file
-        jar_path = ''
         tools_jar_url = self.wrapper_options.get('toolsJar')
         try:
             if tools_jar_url is None:

user_tools/src/spark_rapids_pytools/rapids/tool_ctxt.py

@@ -45,6 +45,9 @@ class ToolContext(YAMLPropertiesContainer):
         Utils.resource_path('csp-resources.tgz'),
         Utils.resource_path('csp-resources')
     ]
+    tools_resource_path: ClassVar[List[str]] = [
+        Utils.resource_path('tools-resources')
+    ]
 
     @classmethod
     def are_resources_prepackaged(cls) -> bool:
@@ -87,6 +90,9 @@ def _init_fields(self):
     def get_deploy_mode(self) -> Any:
         return self.platform_opts.get('deployMode')
 
+    def use_local_tools_jar(self) -> bool:
+        return self.get_ctxt('useLocalToolsJar')
+
     def is_fat_wheel_mode(self) -> bool:
         return self.get_ctxt('fatWheelModeEnabled')
 
@@ -142,9 +148,9 @@ def set_local_workdir(self, parent: str):
         self.logger.info('Dependencies are generated locally in local disk as: %s', dep_folder)
         self.logger.info('Local output folder is set as: %s', exec_root_dir)
 
-    def _identify_fat_wheel_jar(self, resource_files: List[str]) -> None:
+    def _identify_tools_wheel_jar(self, resource_files: List[str]) -> None:
         """
-        Identifies the tools JAR file from resource files in fat wheel mode and sets its name in the context.
+        Identifies the tools JAR file from resource files and sets its name in the context.
         :param resource_files: List of resource files to search for the tools JAR file.
         :raises AssertionError: If the number of matching files is not exactly one.
         """
@@ -155,15 +161,28 @@ def _identify_fat_wheel_jar(self, resource_files: List[str]) -> None:
             (f'Expected exactly one tools JAR file, found {len(matched_files)}. '
              'Rebuild the wheel package with the correct tools JAR file.')
         # set the tools JAR file name in the context
-        self.set_ctxt('fatWheelModeJarFileName', FSUtil.get_resource_name(matched_files[0]))
+        self.set_ctxt('useLocalToolsJar', True)
+        self.set_ctxt('toolsJarFileName', FSUtil.get_resource_name(matched_files[0]))
 
     def load_prepackaged_resources(self):
         """
+        Checks for the tools jar and adds it to context
         Checks if the packaging includes the CSP dependencies. If so, it moves the dependencies
         into the tmp folder. This allows the tool to pick the resources from cache folder.
         """
+        for tools_related_files in self.tools_resource_path:
+            # This function uses a regex based comparison to identify the tools jar file
+            # from the tools-resources directory. The jar is pre-packed in the wheel file
+            # and moved to the work directory when user runs the tool.
+            self.logger.info('Checking for tools related files in %s', tools_related_files)
+            if os.path.exists(tools_related_files):
+                FSUtil.copy_resource(tools_related_files, self.get_cache_folder())
+                self._identify_tools_wheel_jar(FSUtil.get_all_files(tools_related_files))
+
         if not self.are_resources_prepackaged():
+            self.logger.info('No prepackaged resources found.')
             return
+
         self.set_ctxt('fatWheelModeEnabled', True)
         self.logger.info(Utils.gen_str_header('Fat Wheel Mode Is Enabled',
                                               ruler='_', line_width=50))
@@ -173,12 +192,10 @@ def load_prepackaged_resources(self):
                 if os.path.isdir(res_path):
                     # this is a directory, copy all the contents to the tmp
                     FSUtil.copy_resource(res_path, self.get_cache_folder())
-                    self._identify_fat_wheel_jar(FSUtil.get_all_files(res_path))
                 else:
                     # this is an archived file
                     with tarfile.open(res_path, mode='r:*') as tar_file:
                         tar_file.extractall(self.get_cache_folder())
-                        self._identify_fat_wheel_jar(tar_file.getnames())
                         tar_file.close()
 
     def get_output_folder(self) -> str:
@@ -193,11 +210,10 @@ def get_local_work_dir(self) -> str:
         return self.get_local('depFolder')
 
     def get_rapids_jar_url(self) -> str:
-        self.logger.info('Fetching the Rapids Jar URL')
-        # get the version from the package, instead of the yaml file
-        # jar_version = self.get_value('sparkRapids', 'version')
-        if self.is_fat_wheel_mode():
-            return self._get_tools_jar_in_fat_wheel_mode()
+        self.logger.info('Fetching the Rapids Jar URL from local context')
+        if self.use_local_tools_jar():
+            return self._get_tools_jar_from_local()
+        self.logger.info('Tools JAR not found in local context. Downloading from Maven.')
         mvn_base_url = self.get_value('sparkRapids', 'mvnUrl')
         jar_version = Utilities.get_latest_mvn_jar_from_metadata(mvn_base_url)
         rapids_url = self.get_value('sparkRapids', 'repoUrl').format(mvn_base_url, jar_version, jar_version)
@@ -237,22 +253,22 @@ def get_platform_name(self) -> str:
         """
         return CspEnv.pretty_print(self.platform.type_id)
 
-    def _get_tools_jar_in_fat_wheel_mode(self) -> str:
+    def _get_tools_jar_from_local(self) -> str:
         """
         Extracts the tools JAR file from the context and returns its path from the cache folder.
         """
-        jar_filename = self.get_ctxt('fatWheelModeJarFileName')
+        jar_filename = self.get_ctxt('toolsJarFileName')
         if jar_filename is None:
             raise ValueError(
-                'In Fat Mode. Tools JAR file name not found in context. '
-                'Rebuild the wheel package or re-run without fat wheel mode.'
+                'Tools JAR file name not found in context. '
+                'Make sure the tools JAR is included in the package.'
             )
         # construct the path to the tools JAR file in the cache folder
         jar_filepath = FSUtil.build_path(self.get_cache_folder(), jar_filename)
         if not FSUtil.resource_exists(jar_filepath):
             raise FileNotFoundError(
-                f'In Fat Mode. Tools JAR not found in cache folder: {jar_filepath}. '
-                'Rebuild the wheel package or re-run without fat wheel mode.'
+                f'Tools JAR not found in cache folder: {jar_filepath}. '
+                'Rebuild the wheel package'
             )
         self.logger.info('Using jar from wheel file %s', jar_filepath)
         return jar_filepath