Updated download artifact path (#490)

Author: lu-ohai

ads/config.py

@@ -45,6 +45,7 @@
     NB_SESSION_OCID or JOB_RUN_OCID or MD_OCID or PIPELINE_RUN_OCID or DATAFLOW_RUN_OCID
 )
 NO_CONTAINER = os.environ.get("NO_CONTAINER")
+TMPDIR = os.environ.get("TMPDIR")
 
 
 def export(

ads/model/artifact_downloader.py

@@ -63,6 +63,7 @@ def download(self):
                     "Set `force_overwrite` to `True` if you wish to overwrite."
                 )
             shutil.rmtree(self.target_dir)
+            os.mkdir(self.target_dir)
         with utils.get_progress_bar(
             ArtifactDownloader.PROGRESS_STEPS_COUNT + self.PROGRESS_STEPS_COUNT
         ) as progress:
@@ -85,15 +86,16 @@ def _download(self):
         """Downloads model artifacts."""
         self.progress.update("Importing model artifacts from catalog")
         zip_content = self.dsc_model.get_model_artifact_content()
-        with tempfile.TemporaryDirectory() as temp_dir:
-            self.progress.update("Copying model artifacts to the artifact directory")
-            zip_file_path = os.path.join(temp_dir, f"{str(uuid.uuid4())}.zip")
-            with open(zip_file_path, "wb") as zip_file:
-                zip_file.write(zip_content)
-            self.progress.update("Extracting model artifacts")
-            with ZipFile(zip_file_path) as zip_file:
-                zip_file.extractall(self.target_dir)
-
+        self.progress.update("Copying model artifacts to the artifact directory")
+        
+        zip_file_path = os.path.join(self.target_dir, f"{str(uuid.uuid4())}.zip")
+        with open(zip_file_path, "wb") as zip_file:
+            zip_file.write(zip_content)
+        self.progress.update("Extracting model artifacts")
+        with ZipFile(zip_file_path) as zip_file:
+            zip_file.extractall(self.target_dir)
+        
+        utils.remove_file(zip_file_path)
 
 class LargeArtifactDownloader(ArtifactDownloader):
     PROGRESS_STEPS_COUNT = 4
@@ -157,17 +159,19 @@ def _download(self):
 
         self.dsc_model.import_model_artifact(bucket_uri=bucket_uri, region=self.region)
         self.progress.update("Copying model artifacts to the artifact directory")
-        with tempfile.TemporaryDirectory() as temp_dir:
-            zip_file_path = os.path.join(temp_dir, f"{str(uuid.uuid4())}.zip")
-            zip_file_path = utils.copy_file(
-                uri_src=bucket_uri,
-                uri_dst=zip_file_path,
-                auth=self.auth,
-                progressbar_description="Copying model artifacts to the artifact directory",
-            )
-            self.progress.update("Extracting model artifacts")
-            with ZipFile(zip_file_path) as zip_file:
-                zip_file.extractall(self.target_dir)
+        
+        zip_file_path = os.path.join(self.target_dir, f"{str(uuid.uuid4())}.zip")
+        zip_file_path = utils.copy_file(
+            uri_src=bucket_uri,
+            uri_dst=zip_file_path,
+            auth=self.auth,
+            progressbar_description="Copying model artifacts to the artifact directory",
+        )
+        self.progress.update("Extracting model artifacts")
+        with ZipFile(zip_file_path) as zip_file:
+            zip_file.extractall(self.target_dir)
+
+        utils.remove_file(zip_file_path)
 
         if self.remove_existing_artifact:
             self.progress.update(

ads/model/generic_model.py

@@ -31,6 +31,7 @@
     NB_SESSION_OCID,
     PIPELINE_RUN_COMPARTMENT_OCID,
     PROJECT_OCID,
+    TMPDIR,
 )
 from ads.evaluations import EvaluatorMixin
 from ads.feature_engineering import ADSImage
@@ -183,7 +184,7 @@ def _prepare_artifact_dir(artifact_dir: str = None) -> str:
     if artifact_dir and isinstance(artifact_dir, str):
         return os.path.abspath(os.path.expanduser(artifact_dir))
 
-    artifact_dir = tempfile.mkdtemp()
+    artifact_dir = TMPDIR or tempfile.mkdtemp()
     logger.info(
         f"The `artifact_dir` was not provided and "
         f"automatically set to: {artifact_dir}"

tests/unitary/default_setup/model/test_artifact_downloader.py

@@ -100,20 +100,19 @@ def test_downaload_small_artifact(self):
             expected_artifact_bytes_content
         )
         with tempfile.TemporaryDirectory() as tmp_dir:
-            target_dir = os.path.join(tmp_dir, "model_artifacts/")
             SmallArtifactDownloader(
                 dsc_model=self.mock_dsc_model,
-                target_dir=target_dir,
+                target_dir=tmp_dir,
                 force_overwrite=True,
             ).download()
 
             self.mock_dsc_model.get_model_artifact_content.assert_called()
 
             test_files = list(
-                glob.iglob(os.path.join(target_dir, "**"), recursive=True)
+                glob.iglob(os.path.join(tmp_dir, "**"), recursive=True)
             )
             expected_files = [
-                os.path.join(target_dir, file_name)
+                os.path.join(tmp_dir, file_name)
                 for file_name in ["", "runtime.yaml", "score.py"]
             ]
             assert sorted(test_files) == sorted(expected_files)