add secrets to lepton (#383)

* adding secret variable in Lepton

Signed-off-by: Zoey Zhang <[email protected]>

* added tests for env var

Signed-off-by: Zoey Zhang <[email protected]>

* linting

Signed-off-by: Zoey Zhang <[email protected]>

---------

Signed-off-by: Zoey Zhang <[email protected]>

Author: zoeyz101

nemo_run/core/execution/lepton.py

@@ -32,6 +32,7 @@
 from leptonai.api.v1.types.dedicated_node_group import DedicatedNodeGroup
 from leptonai.api.v1.types.deployment import (
     EnvVar,
+    EnvValue,
     LeptonContainer,
     Mount,
 )
@@ -72,6 +73,7 @@ class LeptonExecutor(Executor):
     resource_shape: str = ""
     node_group: str = ""
     node_reservation: str = ""
+    secret_vars: dict[str, str] = field(default_factory=dict)
     mounts: list[dict[str, Any]] = field(default_factory=list)
     lepton_job_dir: str = field(init=False, default="")
     image_pull_secrets: list[str] = field(
@@ -248,6 +250,8 @@ def create_lepton_job(self, name: str):
         client = APIClient()
 
         envs = [EnvVar(name=key, value=value) for key, value in self.env_vars.items()]
+        for key, value in self.secret_vars.items():
+            envs.append(EnvVar(name=key, value_from=EnvValue(secret_name_ref=value)))
 
         cmd = ["/bin/bash", "-c", f"bash {self.lepton_job_dir}/launch_script.sh"]
 

test/core/execution/test_lepton.py

@@ -24,6 +24,8 @@
     LeptonContainer,
     LeptonResourceAffinity,
     Mount,
+    EnvVar,
+    EnvValue,
 )
 from leptonai.api.v1.types.job import LeptonJob, LeptonJobUserSpec
 
@@ -95,6 +97,21 @@ def test_init_without_node_reservation(self):
 
         assert executor.node_reservation == ""
 
+    def test_init_with_secret_vars(self):
+        """Test initialization with node_reservation parameter."""
+        executor = LeptonExecutor(
+            resource_shape="gpu.8xh100-80gb",
+            node_group="my-node-group",
+            container_image="test-image",
+            nodes=2,
+            gpus_per_node=8,
+            secret_vars={"WANDB_API_KEY": "WANDB_API_KEY.zozhang"},
+            nemo_run_dir="/workspace/nemo_run",
+            mounts=[{"path": "/workspace", "mount_path": "/workspace"}],
+        )
+
+        assert executor.secret_vars == {"WANDB_API_KEY": "WANDB_API_KEY.zozhang"}
+
     @patch("nemo_run.core.execution.lepton.APIClient")
     def test_stop_job(self, mock_APIClient):
         mock_instance = MagicMock()
@@ -371,6 +388,8 @@ def test_create_lepton_job(self, mock_APIClient_class):
             container_image="test-image",
             nemo_run_dir="/test/path",
             node_group="123456",
+            env_vars={"TEST_ENV": "test-value"},
+            secret_vars={"TEST_SECRET": "test-secret"},
             mounts=[{"path": "/test", "mount_path": "/test"}],
         )
         executor._valid_node_ids = MagicMock(return_value=valid_node_ids)
@@ -379,6 +398,11 @@ def test_create_lepton_job(self, mock_APIClient_class):
         executor.create_lepton_job("my-lepton-job")
 
         mock_client.job.create.assert_called_once()
+        created_job = mock_client.job.create.call_args[0][0]
+        assert created_job.spec.envs == [
+            EnvVar(name="TEST_ENV", value="test-value"),
+            EnvVar(name="TEST_SECRET", value_from=EnvValue(secret_name_ref="test-secret")),
+        ]
 
     @patch("nemo_run.core.execution.lepton.APIClient")
     def test_create_lepton_job_with_reservation_config(self, mock_APIClient_class):