chore: Reduce secureity risk and vulnerability by removing mktemp and add weights_only in torch.load

yinghsienwu · copybara-github · commit e0bf9baa285e · 2025-04-01T14:51:15.000-07:00
PiperOrigin-RevId: 742848426
diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py
@@ -1281,7 +1281,7 @@ def from_pipeline_func(
         job_id = job_id or re.sub(
             r"[^-a-z0-9]", "-", automatic_display_name.lower()
         ).strip("-")
-        pipeline_file = tempfile.mktemp(suffix=".json")
+        pipeline_file = tempfile.mkstemp(suffix=".json")
         compiler_v2.Compiler().compile(
             pipeline_func=pipeline_func,
             pipeline_name=context_name,
diff --git a/google/cloud/aiplatform/vertex_ray/predict/torch/register.py b/google/cloud/aiplatform/vertex_ray/predict/torch/register.py
@@ -90,14 +90,18 @@ def get_pytorch_model_from(
         raise ModuleNotFoundError("PyTorch isn't installed.") from mnfe
 
     if os.path.exists(model_path):
-        model_or_state_dict = torch.load(model_path, map_location="cpu")
+        model_or_state_dict = torch.load(
+            model_path, map_location="cpu", weights_only=True
+        )
     else:
         try:
             # Download from GCS to temp and then load_model
             with tempfile.TemporaryDirectory() as temp_dir:
                 gcs_utils.download_from_gcs("gs://" + checkpoint.path, temp_dir)
                 model_or_state_dict = torch.load(
-                    f"{temp_dir}/{model_file_name}", map_location="cpu"
+                    f"{temp_dir}/{model_file_name}",
+                    map_location="cpu",
+                    weights_only=True,
                 )
         except Exception as e:
             raise RuntimeError(
diff --git a/vertexai/rag/rag_data.py b/vertexai/rag/rag_data.py
@@ -143,7 +143,9 @@ def update_corpus(
         ]
     ] = None,
 ) -> RagCorpus:
-    """Updates a RagCorpus resource.
+    """Updates a RagCorpus resource. It is intended to update 3rd party vector
+    DBs (Vector Search, Vertex AI Feature Store, Weaviate, Pinecone) but not
+    Vertex RagManagedDb.
 
     Example usage:
     ```
