JobDefinition env parameter to set pip index url and host

Author: ChristianZaccaria

Diff for: src/codeflare_sdk/job/jobs.py

@@ -22,6 +22,8 @@
 from torchx.schedulers.ray_scheduler import RayScheduler
 from torchx.specs import AppHandle, parse_app_handle, AppDryRunInfo
 
+from ..utils.generate_yaml import update_pip_requirements
+
 
 if TYPE_CHECKING:
     from ..cluster.cluster import Cluster
@@ -90,6 +92,12 @@ def __init__(
         )
         self.image = image
         self.workspace = workspace
+        if 'PIP_INDEX_URL' in self.env or 'PIP_TRUSTED_HOST' in self.env:
+            update_pip_requirements(self)
+        else:
+            self.env.setdefault('PIP_INDEX_URL', 'https://pypi.org/simple')
+            self.env.setdefault('PIP_TRUSTED_HOST', 'pypi.org')
+            update_pip_requirements(self)
 
     def _dry_run(self, cluster: "Cluster"):
         j = f"{cluster.config.num_workers}x{max(cluster.config.num_gpus, 1)}"  # # of proc. = # of gpus

Diff for: src/codeflare_sdk/utils/generate_yaml.py

@@ -21,6 +21,7 @@
 import sys
 import os
 import argparse
+from pathlib import Path
 import uuid
 from kubernetes import client, config
 from .kube_api_helpers import _kube_api_error_handling
@@ -607,7 +608,6 @@ def write_components(user_yaml: dict, output_file_name: str):
                 )
     print(f"Written to: {output_file_name}")
 
-
 def generate_appwrapper(
     name: str,
     namespace: str,
@@ -689,3 +689,37 @@ def generate_appwrapper(
     else:
         write_user_appwrapper(user_yaml, outfile)
     return outfile
+
+def update_pip_requirements(self):
+    pip_index_url = self.env.get('PIP_INDEX_URL')
+    pip_trusted_host = self.env.get('PIP_TRUSTED_HOST')
+    requirements_path = Path('requirements.txt')
+
+    if requirements_path.exists():
+        with requirements_path.open('r') as file:
+            requirements = file.readlines()
+
+        # Check and replace or add --trusted-host and --index-url
+        trusted_host = f"--trusted-host {pip_trusted_host}\n"
+        index_url = f"--index-url {pip_index_url}\n"
+        modified_requirements = []
+
+        for line in requirements:
+            if line.startswith("--trusted-host"):
+                modified_requirements.append(trusted_host)
+                trusted_host = None
+            elif line.startswith("--index-url"):
+                modified_requirements.append(index_url)
+                index_url = None
+            else:
+                modified_requirements.append(line)
+
+        # Append the lines if they were not replaced
+        if index_url:
+            modified_requirements.insert(0, index_url)
+        if trusted_host:
+            modified_requirements.insert(0, trusted_host)
+
+        # Write back the modified requirements
+        with requirements_path.open('w') as file:
+            file.writelines(modified_requirements)

Diff for: tests/unit_test.py

@@ -2091,7 +2091,7 @@ def test_DDPJobDefinition_creation():
     assert ddp.memMB == 1024
     assert ddp.h == None
     assert ddp.j == "2x1"
-    assert ddp.env == {"test": "test"}
+    assert ddp.env == {"PIP_INDEX_URL":"https://pypi.org/simple", "PIP_TRUSTED_HOST": "pypi.org", "test": "test"}
     assert ddp.max_retries == 0
     assert ddp.mounts == []
     assert ddp.rdzv_port == 29500