Print venv in launch

Author: XkunW

vec_inf/cli/_helper.py

@@ -76,6 +76,7 @@ def format_table_output(self) -> Table:
         table.add_row("Memory/Node", self.params["mem_per_node"])
 
         # Add job config details
+        table.add_row("Virtual Environment", self.params["venv"])
         table.add_row(
             "Model Weights Directory",
             str(Path(self.params["model_weights_parent_dir"], self.model_name)),

vec_inf/client/_helper.py

@@ -31,6 +31,7 @@
     BatchSlurmScriptGenerator,
     SlurmScriptGenerator,
 )
+from vec_inf.client._slurm_vars import CONTAINER_MODULE_NAME, IMAGE_PATH
 from vec_inf.client.config import ModelConfig
 from vec_inf.client.models import (
     BatchLaunchResponse,
@@ -332,6 +333,10 @@ def launch(self) -> LaunchResponse:
             job_log_dir / f"{self.model_name}.{self.slurm_job_id}.sbatch"
         )
 
+        # Replace venv with image path if using container
+        if self.params["venv"] == CONTAINER_MODULE_NAME:
+            self.params["venv"] = IMAGE_PATH
+
         with job_json.open("w") as file:
             json.dump(self.params, file, indent=4)