From dd36593eecb28a1d7bbed1ed839c295ee7042a23 Mon Sep 17 00:00:00 2001
From: Matteo Bunino <matteo.bunino@gmail.com>
Date: Tue, 5 Nov 2024 14:15:51 +0100
Subject: [PATCH] Add chaining

---
 ci/run.sh               |  4 ++++
 ci/src/main/__init__.py | 41 ++++++++---------------------------------
 2 files changed, 12 insertions(+), 33 deletions(-)
 create mode 100644 ci/run.sh

diff --git a/ci/run.sh b/ci/run.sh
new file mode 100644
index 00000000..38f117e3
--- /dev/null
+++ b/ci/run.sh
@@ -0,0 +1,4 @@
+# Example of running dagger pipeline of build and local test for torch container
+dagger call \
+    build-torch --context=.. --dockerfile=../env-files/torch/Dockerfile \
+    test-torch
\ No newline at end of file
diff --git a/ci/src/main/__init__.py b/ci/src/main/__init__.py
index 2ed68443..2e3e04cb 100644
--- a/ci/src/main/__init__.py
+++ b/ci/src/main/__init__.py
@@ -21,7 +21,8 @@
 # For example, to import from src/main/main.py:
 # >>> from .main import Itwinai as Itwinai
 
-from typing import Annotated
+from typing import Annotated, Optional, Self
+import dataclasses
 
 import dagger
 from dagger import dag, function, object_type, Doc
@@ -30,7 +31,7 @@
 @object_type
 class Itwinai:
 
-    # torch_container: dagger.Container = None
+    torch_container: Optional[dagger.Container] = dataclasses.field(default=None, init=False)
 
     @function
     def build_torch(
@@ -43,25 +44,16 @@ def build_torch(
             str,
             Doc("location of Dockerfile"),
         ],
-    ) -> dagger.Container:
+    ) -> Self:
         """Build itwinai torch container image from existing Dockerfile"""
-        return (
+        self.torch_container = (
             dag.container()
             .build(context=context, dockerfile=dockerfile)
         )
+        return self
 
     @function
-    async def test_torch(
-        self,
-        context: Annotated[
-            dagger.Directory,
-            Doc("location of source directory"),
-        ],
-        dockerfile: Annotated[
-            str,
-            Doc("location of Dockerfile"),
-        ],
-    ) -> str:
+    async def test_torch(self) -> str:
         """Test itwinai torch container image with pytest on non-HPC environments."""
         test_cmd = [
             "pytest",
@@ -71,24 +63,7 @@ async def test_torch(
             "tests"
         ]
         return await (
-            self.build_torch(context=context, dockerfile=dockerfile)
+            self.torch_container
             .with_exec(test_cmd)
             .stdout()
         )
-
-    @function
-    def container_echo(self, string_arg: str) -> dagger.Container:
-        """Returns a container that echoes whatever string argument is provided"""
-        return dag.container().from_("alpine:latest").with_exec(["echo", string_arg])
-
-    @function
-    async def grep_dir(self, directory_arg: dagger.Directory, pattern: str) -> str:
-        """Returns lines that match a pattern in the files of the provided Directory"""
-        return await (
-            dag.container()
-            .from_("alpine:latest")
-            .with_mounted_directory("/mnt", directory_arg)
-            .with_workdir("/mnt")
-            .with_exec(["grep", "-R", pattern, "."])
-            .stdout()
-        )