Enable Spark Fast Register

flytekit/core/tracker.py

@@ -270,6 +270,9 @@ def _resolve_abs_module_name(self, path: str, package_root: typing.Optional[str]
         # Let us remove any extensions like .py
         basename = os.path.splitext(basename)[0]
 
+        if not Path(dirname).is_dir():
+            return basename
+
         if dirname == package_root:
             return basename
 

plugins/flytekit-spark/flytekitplugins/spark/task.py

@@ -1,8 +1,11 @@
+import glob
 import os
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, Optional, Union, cast
 
+import time
 import click
+import shutil
 from google.protobuf.json_format import MessageToDict
 
 from flytekit import FlyteContextManager, PythonFunctionTask, lazy_module, logger
@@ -158,13 +161,6 @@ def __init__(
             **kwargs,
         )
 
-    def get_image(self, settings: SerializationSettings) -> str:
-        if isinstance(self.container_image, ImageSpec):
-            # Ensure that the code is always copied into the image, even during fast-registration.
-            self.container_image.source_root = settings.source_root
-
-        return get_registerable_container_image(self.container_image, settings.image_config)
-
     def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
         job = SparkJob(
             spark_conf=self.task_config.spark_conf,
@@ -201,6 +197,20 @@ def pre_execute(self, user_params: ExecutionParameters) -> ExecutionParameters:
             sess_builder = sess_builder.config(conf=spark_conf)
 
         self.sess = sess_builder.getOrCreate()
+
+        current_time = time.time()
+        current_dir = os.getcwd()
+        # Set the modification time of all files in the current directory to the current time
+        # since fast register doens't preserve the modification time of the files and make_archive
+        # does not support timestamps before 1980
+        for foldername, subfolders, filenames in os.walk(current_dir):
+            for filename in filenames:
+                file_path = os.path.join(foldername, filename)
+                os.utime(file_path, (current_time, current_time))
+
+        shutil.make_archive("archive", 'zip', current_dir)
+        self.sess.sparkContext.addPyFile("archive.zip")
+
         return user_params.builder().add_attr("SPARK_SESSION", self.sess).build()
 
     def execute(self, **kwargs) -> Any: