Merge pull request #36 from higgsfield-ai/dev

merge: dev

higgsfield/internal/ci/cli.py

@@ -54,7 +54,8 @@ def decode_secrets(env: str):
 
 
 @click.command("setup-nodes")
-def setup_nodes():
+@click.option('--invoker_tag', default="v0.0.1", help="Tag of the invoker binary to use")
+def setup_nodes(invoker_tag: str = "v0.0.1"):
     wd = wd_path()
     app_config = AppConfig.from_path(wd)
 
@@ -74,7 +75,7 @@ def setup_nodes():
 
     app_config.set_git_origin_url(project_path)
 
-    setup = Setup(app_config, project_path)
+    setup = Setup(app_config, project_path, invoker_tag)
 
     try:
         setup.create_ssh_key_file()

higgsfield/internal/ci/setup.py

@@ -11,10 +11,12 @@ from higgsfield.internal.experiment.builder import header
 
 
 docker_install_script = '''/bin/bash -c "$(curl -fsSL https://gist.githubusercontent.com/arpanetus/1c1210b9e432a04dcfb494725a407a70/raw/5d47baa19b7100261a2368a43ace610528e0dfa2/install.sh)"'''
-invoker_install_script = """wget https://github.com/ml-doom/invoker/releases/download/latest/invoker-latest-linux-amd64.tar.gz && \
-        tar -xvf invoker-latest-linux-amd64.tar.gz && \
+
+def invoker_install_script(tag: str) -> str:
+    return f"""wget https://github.com/ml-doom/invoker/releases/download/{tag}/invoker-{tag}-linux-amd64.tar.gz && \
+        tar -xvf invoker-{tag}-linux-amd64.tar.gz && \
         sudo mv invoker /usr/bin/invoker && \
-        rm invoker-latest-linux-amd64.tar.gz"""
+        rm invoker-{tag}-linux-amd64.tar.gz"""
 
 
 def deploy_key_script(key: str, project_name: str, deploy_key_string: str):
@@ -32,11 +34,13 @@ class Setup:
     path: str
     deploy_key: str
     project_path: Path
+    invoker_tag: str
 
     def __init__(
         self,
         app_config: AppConfig,
         project_path: Path,
+        invoker_tag: str,
     ):
         self.app_config = app_config
 
@@ -44,6 +48,7 @@ class Setup:
             raise ValueError(reason)
 
         self.project_path = project_path
+        self.invoker_tag = invoker_tag
 
     def create_ssh_key_file(self):
         if self.app_config.key is None:
@@ -97,7 +102,7 @@ class Setup:
         print("\n\n\nINSTALLING INVOKER\n\n\n")
         to_run = []
         for conn in self.connections:
-            to_run.append(printer(await conn.run(invoker_install_script)))
+            to_run.append(printer(await conn.run(invoker_install_script(self.invoker_tag))))
 
         await asyncio.gather(*to_run)
 
@@ -111,6 +116,15 @@ class Setup:
             to_run.append(printer(await conn.run(dk_script)))
     
         await asyncio.gather(*to_run)
+
+        # close connections
+        for conn in self.connections:
+            conn.close()
+
+
+        # re-establish connections
+        await self.establish_connections()
+
         
         print("\n\n\nPULLING BASE DOCKER IMAGE\n\n\n")
         to_run = []

higgsfield/internal/experiment/params.py

@@ -103,7 +103,7 @@ class Param:
                 f"{indent}description: {remove_trailing_yaml(safe_dump(self.description))}"
             )
         to_join.append(f"{indent}required: {self.required}")
-        if self.default:
+        if self.default is not None:
             d = self.default
             if type(self.default) == str:
                 d = remove_trailing_yaml(safe_dump(self.default))