add runtime dockerize ut

client/tests/core/test_runtime.py

@@ -24,7 +24,12 @@
 from starwhale.utils.error import UnExpectedConfigFieldError
 from starwhale.utils.config import SWCliConfigMixed
 from starwhale.core.runtime.view import RuntimeTermView
-from starwhale.core.runtime.model import Runtime, RuntimeConfig, StandaloneRuntime
+from starwhale.core.runtime.model import (
+    Runtime,
+    _TEMPLATE_DIR,
+    RuntimeConfig,
+    StandaloneRuntime,
+)
 
 
 class StandaloneRuntimeTestCase(TestCase):
@@ -134,6 +139,72 @@ def test_quickstart_from_ishell_conda(self, m_call: MagicMock) -> None:
             )
         )
 
+    @patch("starwhale.core.runtime.model.StandaloneRuntime.restore")
+    @patch("starwhale.base.bundle.extract_tar")
+    @patch("starwhale.core.runtime.model.BundleCopy")
+    def test_quickstart_from_uri(
+        self, m_bundle_copy: MagicMock, m_extract: MagicMock, m_restore: MagicMock
+    ) -> None:
+        workdir = Path("/home/starwhale/myproject")
+        name = "rttest"
+        version = "112233"
+        cloud_uri = URI(f"http://0.0.0.0:80/project/1/runtime/{name}/version/{version}")
+        ensure_dir(workdir / ".extract")
+
+        runtime_config = self.get_runtime_config()
+        runtime_config["name"] = name
+        extract_dir = workdir / ".extract"
+        ensure_file(
+            extract_dir / DefaultYAMLName.RUNTIME,
+            content=yaml.safe_dump(runtime_config),
+        )
+        ensure_dir(extract_dir / "wheels")
+        ensure_dir(extract_dir / "dependencies")
+        ensure_file(extract_dir / "wheels" / "dummy.whl", "")
+        ensure_file(extract_dir / "dependencies" / "requirements.txt", "numpy")
+        ensure_file(
+            extract_dir / "_manifest.yaml",
+            yaml.safe_dump({"environment": {"mode": "venv"}}),
+        )
+
+        sw = SWCliConfigMixed()
+        runtime_dir = os.path.join(sw.rootdir, "self", "runtime")
+        bundle_path = os.path.join(
+            runtime_dir, name, f"{version[:VERSION_PREFIX_CNT]}", f"{version}.swrt"
+        )
+        ensure_dir(os.path.dirname(bundle_path))
+        ensure_file(bundle_path, "")
+        venv_dir = workdir / ".venv"
+        ensure_dir(venv_dir)
+
+        assert not (workdir / "dummy.whl").exists()
+        assert not (venv_dir / ".gitignore").exists()
+        assert not (workdir / "requirements.txt").exists()
+
+        StandaloneRuntime.quickstart_from_uri(
+            workdir=workdir,
+            name=name,
+            uri=cloud_uri,
+            force=True,
+            restore=True,
+        )
+
+        runtime_path = workdir / DefaultYAMLName.RUNTIME
+        assert runtime_path.exists()
+        runtime_config = yaml.safe_load(runtime_path.read_text())
+        assert runtime_config["name"] == name
+        assert runtime_config["mode"] == "venv"
+        assert runtime_config["dependencies"][0] == "requirements.txt"
+        assert runtime_config["dependencies"][1]["wheels"] == ["dummy.whl"]
+        assert runtime_config["dependencies"][2]["pip"] == ["Pillow"]
+        assert (workdir / "dummy.whl").exists()
+        assert (venv_dir / ".gitignore").exists()
+        assert (workdir / "requirements.txt").exists()
+
+        assert m_bundle_copy.call_count == 1
+        assert m_extract.call_count == 1
+        assert m_restore.call_args[0] == (extract_dir, venv_dir)
+
     @patch("starwhale.utils.venv.get_user_runtime_python_bin")
     @patch("starwhale.utils.venv.check_call")
     @patch("starwhale.utils.venv.subprocess.check_output", return_value=b"3.7")
@@ -151,13 +222,28 @@ def test_build_venv(
         runtime_config = self.get_runtime_config()
         runtime_config["environment"]["cuda"] = "11.5"
         runtime_config["environment"]["cudnn"] = "8"
+        runtime_config["dependencies"].append(
+            {
+                "files": [
+                    {
+                        "dest": "bin/prepare.sh",
+                        "name": "prepare",
+                        "post": "bash bin/prepare.sh",
+                        "pre": "ls bin/prepare.sh",
+                        "src": "prepare.sh",
+                    }
+                ]
+            }
+        )
         self.fs.create_file(
             os.path.join(workdir, DefaultYAMLName.RUNTIME),
             contents=yaml.safe_dump(runtime_config),
         )
+        self.fs.create_file(os.path.join(workdir, "prepare.sh"), contents="")
         self.fs.create_file(
             os.path.join(workdir, "requirements.txt"), contents="requests==2.0.0"
         )
+        self.fs.create_file(os.path.join(workdir, "dummy.whl"), contents="")
 
         uri = URI(name, expected_type=URIType.RUNTIME)
         sr = StandaloneRuntime(uri)
@@ -185,6 +271,9 @@ def test_build_venv(
 
         assert os.path.exists(bundle_path)
         assert os.path.exists(runtime_workdir)
+        assert os.path.exists(os.path.join(runtime_workdir, "wheels", "dummy.whl"))
+        assert os.path.exists(os.path.join(runtime_workdir, "files/bin/prepare.sh"))
+
         assert "latest" in sr.tag.list()
 
         _manifest = load_yaml(os.path.join(runtime_workdir, DEFAULT_MANIFEST_NAME))
@@ -207,6 +296,15 @@ def test_build_venv(
         assert _manifest["version"] == sr.uri.object.version
         assert _manifest["environment"]["mode"] == "venv"
         assert _manifest["environment"]["lock"]["shell"]["use_venv"]
+        assert _manifest["artifacts"]["wheels"] == ["wheels/dummy.whl"]
+        assert _manifest["artifacts"]["files"][0] == {
+            "_swrt_dest": "files/bin/prepare.sh",
+            "dest": "bin/prepare.sh",
+            "name": "prepare",
+            "post": "bash bin/prepare.sh",
+            "pre": "ls bin/prepare.sh",
+            "src": "prepare.sh",
+        }
         assert not _manifest["dependencies"]["local_packaged_env"]
 
         uri = URI(name, expected_type=URIType.RUNTIME)
@@ -316,11 +414,16 @@ def test_build_conda(
 
         runtime_config = self.get_runtime_config()
         runtime_config["mode"] = "conda"
+        runtime_config["dependencies"].append("conda.yaml")
+        runtime_config["dependencies"].append("unparsed.xxx")
         self.fs.create_file(
             os.path.join(workdir, DefaultYAMLName.RUNTIME),
             contents=yaml.safe_dump(runtime_config),
         )
         self.fs.create_file(os.path.join(workdir, "requirements.txt"), contents="")
+        self.fs.create_file(os.path.join(workdir, "conda.yaml"), contents="")
+        self.fs.create_file(os.path.join(workdir, "unparsed.xxx"), contents="")
+        self.fs.create_file(os.path.join(workdir, "dummy.whl"), contents="")
         uri = URI(name, expected_type=URIType.RUNTIME)
         sr = StandaloneRuntime(uri)
         sr.build(Path(workdir))
@@ -336,6 +439,8 @@ def get_runtime_config(self) -> t.Dict[str, t.Any]:
             },
             "dependencies": [
                 "requirements.txt",
+                {"wheels": ["dummy.whl"]},
+                {"pip": ["Pillow"]},
             ],
         }
 
@@ -349,8 +454,15 @@ def test_restore_venv(
         export_dir = os.path.join(workdir, "export")
         venv_dir = os.path.join(export_dir, "venv")
         dep_dir = os.path.join(workdir, "dependencies")
+        scripts_dir = os.path.join(workdir, "files", "bin")
+        wheels_dir = os.path.join(workdir, "wheels")
 
         ensure_dir(workdir)
+        ensure_dir(scripts_dir)
+        ensure_dir(wheels_dir)
+        self.fs.create_file(os.path.join(scripts_dir, "prepare.sh"), contents="")
+        self.fs.create_file(os.path.join(wheels_dir, "dummy.wheel"), contents="")
+
         self.fs.create_file(
             os.path.join(workdir, DEFAULT_MANIFEST_NAME),
             contents=yaml.safe_dump(
@@ -367,6 +479,19 @@ def test_restore_venv(
                             "requirements-test.txt",
                         ],
                     },
+                    "artifacts": {
+                        "files": [
+                            {
+                                "_swrt_dest": "files/bin/prepare.sh",
+                                "dest": "bin/prepare.sh",
+                                "name": "prepare",
+                                "post": "bash bin/prepare.sh",
+                                "pre": "ls bin/prepare.sh",
+                                "src": "scripts/prepare.sh",
+                            }
+                        ],
+                        "wheels": ["wheels/dummy.whl"],
+                    },
                 }
             ),
         )
@@ -399,6 +524,7 @@ def test_restore_venv(
             "--pre",
             "starwhale",
         ]
+        assert (Path(workdir) / "export/venv/bin/prepare.sh").exists()
 
         m_call.reset_mock()
         m_exists.return_value = True
@@ -575,3 +701,160 @@ def test_lock_conda(self, m_output: MagicMock, m_call: MagicMock) -> None:
             "/tmp/conda",
             "--file",
         ]
+
+    def get_mock_manifest(self) -> t.Dict[str, t.Any]:
+        return {
+            "name": "rttest",
+            "version": "112233",
+            "base_image": "ghcr.io/star-whale/starwhale:latest-cuda11.4",
+            "dependencies": {
+                "conda_files": [],
+                "conda_pkgs": [],
+                "pip_pkgs": ["numpy"],
+                "pip_files": ["requirements-sw-lock.txt"],
+                "local_packaged_env": False,
+            },
+            "configs": {
+                "conda": {"channels": ["conda-forge"]},
+                "docker": {"image": "ghcr.io/star-whale/runtime/pytorch"},
+                "pip": {
+                    "extra_index_url": ["https://pypi.doubanio.com/simple"],
+                    "index_url": "https://pypi.tuna.tsinghua.edu.cn/simple",
+                    "trusted_host": ["pypi.tuna.tsinghua.edu.cn", "pypi.doubanio.com"],
+                },
+            },
+            "artifacts": {
+                "dependencies": ["dependencies/requirements-sw-lock.txt"],
+                "files": [
+                    {
+                        "_swrt_dest": "files/bin/prepare.sh",
+                        "dest": "bin/prepare.sh",
+                        "name": "prepare",
+                        "post": "bash bin/prepare.sh",
+                        "pre": "ls bin/prepare.sh",
+                        "src": "scripts/prepare.sh",
+                    }
+                ],
+                "runtime_yaml": "runtime.yaml",
+                "wheels": ["wheels/dummy-0.0.0-py3-none-any.whl"],
+            },
+            "environment": {
+                "arch": ["noarch"],
+                "auto_lock_dependencies": False,
+                "lock": {
+                    "env_name": "",
+                    "env_prefix_path": "",
+                    "shell": {
+                        "python_env": "conda",
+                        "python_version": "3.8.13",
+                        "use_conda": True,
+                        "use_venv": False,
+                    },
+                    "starwhale_version": "0.0.0.dev0",
+                    "system": "Linux",
+                    "use_shell_detection": True,
+                },
+                "mode": "venv",
+                "python": "3.8",
+            },
+        }
+
+    @patch("starwhale.utils.docker.check_call")
+    def test_dockerize(self, m_check: MagicMock) -> None:
+        self.fs.add_real_directory(_TEMPLATE_DIR)
+        name = "rttest"
+        version = "112233"
+        image = "docker.io/t1/t2"
+        uri = URI(f"{name}/version/{version}", expected_type=URIType.RUNTIME)
+        manifest = self.get_mock_manifest()
+        manifest["name"] = name
+        manifest["version"] = version
+        manifest["configs"]["docker"]["image"] = image
+
+        sr = StandaloneRuntime(uri)
+        ensure_dir(sr.store.snapshot_workdir)
+        ensure_file(sr.store.manifest_path, content=yaml.safe_dump(manifest))
+        sr.dockerize(
+            tags=["t1", "t2"],
+            platforms=[SupportArch.AMD64],
+            push=True,
+            dry_run=False,
+            use_starwhale_builder=True,
+            reset_qemu_static=True,
+        )
+
+        dockerfile_path = sr.store.export_dir / "docker" / "Dockerfile"
+        dockerignore_path = sr.store.snapshot_workdir / ".dockerignore"
+        assert dockerfile_path.exists()
+        assert dockerignore_path.exists()
+        dockerfile_content = dockerfile_path.read_text()
+        assert f"BASE_IMAGE={manifest['base_image']}" in dockerfile_content
+        assert f"starwhale_runtime_version={version}" in dockerfile_content
+
+        assert m_check.call_count == 3
+        assert m_check.call_args_list[0][0][0] == [
+            "docker",
+            "run",
+            "--rm",
+            "--privileged",
+            "multiarch/qemu-user-static",
+            "--reset",
+            "-p",
+            "-yes",
+        ]
+        assert m_check.call_args_list[1][0][0] == [
+            "docker",
+            "buildx",
+            "inspect",
+            "--builder",
+            "starwhale-multiarch-runtime-builder",
+        ]
+
+        build_cmd = " ".join(m_check.call_args_list[2][0][0])
+        assert "--builder starwhale-multiarch-runtime-builder" in build_cmd
+        assert "--platform linux/amd64" in build_cmd
+        assert "--tag t1" in build_cmd
+        assert "--tag t2" in build_cmd
+        assert f"--tag {image}:{version}" in build_cmd
+        assert "--push" in build_cmd
+        assert f"--file {dockerfile_path}" in build_cmd
+
+        RuntimeTermView(f"{name}/version/{version}").dockerize(
+            tags=[],
+            push=False,
+            platforms=[SupportArch.ARM64],
+            dry_run=False,
+            use_starwhale_builder=False,
+            reset_qemu_static=False,
+        )
+
+    @patch("os.execl")
+    def test_activate(self, m_execl: MagicMock) -> None:
+        sw = SWCliConfigMixed()
+        name = "rttest"
+        version = "123"
+        snapshot_dir = (
+            sw.rootdir
+            / "self"
+            / "workdir"
+            / "runtime"
+            / name
+            / f"{version[:VERSION_PREFIX_CNT]}"
+            / version
+        )
+        manifest = self.get_mock_manifest()
+        manifest["name"] = name
+        manifest["version"] = version
+        ensure_dir(snapshot_dir)
+        ensure_file(snapshot_dir / DEFAULT_MANIFEST_NAME, yaml.safe_dump(manifest))
+
+        uri = f"{name}/version/{version}"
+        StandaloneRuntime.activate(uri=uri)
+
+        runtime_config = self.get_runtime_config()
+        runtime_config["mode"] = "conda"
+        ensure_file(
+            snapshot_dir / DefaultYAMLName.RUNTIME, yaml.safe_dump(runtime_config)
+        )
+
+        StandaloneRuntime.activate(path=str(snapshot_dir))