thestage 0.5.43__tar.gz → 0.5.44__tar.gz

{thestage-0.5.43 → thestage-0.5.44}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: thestage
-Version: 0.5.43
+Version: 0.5.44
 Summary: 
 Author: TheStage AI team
 Author-email: hello@thestage.ai

{thestage-0.5.43 → thestage-0.5.44}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "thestage"
-version = "0.5.43"
+version = "0.5.44"
 description = ""
 authors = ["TheStage AI team <hello@thestage.ai>"]
 readme = "README.md"

{thestage-0.5.43 → thestage-0.5.44}/thestage/__init__.py

@@ -1,3 +1,3 @@
 from . import *
 __app_name__ = "thestage"
-__version__ = "0.5.43"
+__version__ = "0.5.44"

{thestage-0.5.43 → thestage-0.5.44}/thestage/color_scheme/color_scheme.py

@@ -1,6 +1,7 @@
 from enum import Enum
 
-
+# https://rich.readthedocs.io/en/stable/appendix/colors.html
 class ColorScheme(str, Enum):
     GIT_HEADLESS = "orange_red1"
     WARNING = "orange_red1"
+    USEFUL_INFO = "deep_sky_blue1"

{thestage-0.5.43 → thestage-0.5.44}/thestage/controllers/project_controller.py

@@ -183,6 +183,34 @@ def run(
     raise typer.Exit(0)
 
 
+@task_app.command(name='cancel', no_args_is_help=True, help=__("Cancel a task by ID"))
+def cancel_task(
+        task_id: Annotated[int, typer.Argument(
+            help=__("Task ID (required)"),
+        )],
+):
+    """
+        Cancels a task
+    """
+    app_logger.info(f'Start cancel task from {get_current_directory()}')
+
+    if not task_id:
+        typer.echo(__('Task ID is required'))
+        raise typer.Exit(1)
+
+    service_factory = validate_config_and_get_service_factory()
+    config = service_factory.get_config_provider().get_full_config()
+
+    project_service = service_factory.get_project_service()
+
+    project_service.cancel_task(
+        config=config,
+        task_id=task_id
+    )
+
+    raise typer.Exit(0)
+
+
 @task_app.command("ls", help=__("List tasks"))
 def list_runs(
         project_uid: Annotated[str, typer.Argument(help=__("Project unique ID. By default, project info is taken from the current directory"), metavar="OPTIONAL")] = None,
@@ -287,9 +315,9 @@ def checkout_project(
         ),
 ):
     """
-        Initializes project in current working directory
+        Checkout project in current working directory
     """
-    app_logger.info(f'Start project init from {get_current_directory()}')
+    app_logger.info(f'Start project checkout from {get_current_directory()}')
 
     service_factory = validate_config_and_get_service_factory(working_directory=working_directory)
     config = service_factory.get_config_provider().get_full_config()
@@ -315,11 +343,8 @@ def checkout_project(
     raise typer.Exit(0)
 
 
-@config_app.command(name='set-default-container', no_args_is_help=True, help=__("Set default docker container for a project installation"))
-def set_default_container(
-        container_uid: Annotated[Optional[str], typer.Argument(
-            help=__("Unique ID of the container to use by default for running tasks"),
-        )] = None,
+@app.command(name='pull', help=__("Pulls the changes from the remote project repository. Equivalent to 'git pull'."))
+def pull_project(
         working_directory: Optional[str] = typer.Option(
             None,
             "--working-directory",
@@ -327,6 +352,56 @@ def set_default_container(
             help=__("Full path to working directory"),
             is_eager=False,
         ),
+):
+    """
+        Pull project in current working directory
+    """
+    app_logger.info(f'Start project pull from {get_current_directory()}')
+
+    service_factory = validate_config_and_get_service_factory(working_directory=working_directory)
+    config = service_factory.get_config_provider().get_full_config()
+
+    project_service = service_factory.get_project_service()
+
+    project_service.pull_project(
+        config=config,
+    )
+
+    raise typer.Exit(0)
+
+
+@app.command(name='reset', help=__("Resets the current project branch to remote counterpart. All working tree changes will be lost. Equivalent to 'git fetch && git reset --hard origin/{ref}'."))
+def reset_project(
+        working_directory: Optional[str] = typer.Option(
+            None,
+            "--working-directory",
+            "-wd",
+            help=__("Full path to working directory"),
+            is_eager=False,
+        ),
+):
+    """
+        Pull project in current working directory
+    """
+    app_logger.info(f'Start project pull from {get_current_directory()}')
+
+    service_factory = validate_config_and_get_service_factory(working_directory=working_directory)
+    config = service_factory.get_config_provider().get_full_config()
+
+    project_service = service_factory.get_project_service()
+
+    project_service.reset_project(
+        config=config
+    )
+
+    raise typer.Exit(0)
+
+
+@config_app.command(name='set-default-container', no_args_is_help=True, help=__("Set default docker container for a project installation"))
+def set_default_container(
+        container_uid: Annotated[Optional[str], typer.Argument(
+            help=__("Unique ID of the container to use by default for running tasks"),
+        )] = None,
         unset_default_container: Optional[bool] = typer.Option(
             False,
             "--unset",
@@ -334,6 +409,13 @@ def set_default_container(
             help=__("Unsets the default docker container"),
             is_eager=False,
         ),
+        working_directory: Optional[str] = typer.Option(
+            None,
+            "--working-directory",
+            "-wd",
+            help=__("Full path to working directory"),
+            is_eager=False,
+        ),
 ):
     """
         Initializes project in current working directory

{thestage-0.5.43 → thestage-0.5.44}/thestage/services/clients/git/git_client.py

@@ -10,6 +10,7 @@ from gitdb.exc import BadName
 from rich import print
 
 from thestage.color_scheme.color_scheme import ColorScheme
+from thestage.config import THESTAGE_CONFIG_DIR
 from thestage.exceptions.git_access_exception import GitAccessException
 from thestage.git.ProgressPrinter import ProgressPrinter
 from thestage.services.filesystem_service import FileSystemServiceCore
@@ -18,7 +19,7 @@ from thestage.services.filesystem_service import FileSystemServiceCore
 class GitLocalClient:
     __base_name_remote: str = 'origin'
     __base_name_local: str = 'main'
-    __git_ignore_thestage_line: str = '/.thestage/'
+    __git_ignore_thestage_line: str = f'/{THESTAGE_CONFIG_DIR}/'
 
     __special_main_branches = ['main', 'master']
 
@@ -30,6 +31,7 @@ class GitLocalClient:
     ):
         self.__file_system_service = file_system_service
 
+    # todo delete this fuckery
     def __get_repo(self, path: str) -> Repo:
         return git.Repo(path)
 
@@ -105,32 +107,37 @@ class GitLocalClient:
         with repo.git.custom_environment(GIT_SSH_COMMAND=git_ssh_cmd):
             remote: Remote = repo.remote(self.__base_name_remote)
             if remote:
+                progress = ProgressPrinter()
                 try:
-                    remote.fetch()
-                except GitCommandError as base_ex:
-                    msg = base_ex.stderr
-                    if msg and 'fatal: Could not read from remote repository' in msg:
-                        raise GitAccessException(
-                            message='You dont have access to repository, or repository not found.',
-                            url=self.build_http_repo_url(git_path=remote.url),
-                            dop_message=msg,
-                        )
-                    else:
-                        raise base_ex
-
-    def git_pull(self, path: str, deploy_key_path: str, branch: Optional[str] = None):
+                    remote.fetch(progress=progress)
+                except GitCommandError as ex:
+                    for line in progress.allDroppedLines():
+                        # returning the whole output if failed - so that user have any idea what's going on
+                        print(f'>> {line}')
+                    raise ex
+
+
+    def git_pull(self, path: str, deploy_key_path: str):
         repo = self.__get_repo(path=path)
         git_ssh_cmd = 'ssh -F /dev/null -o StrictHostKeyChecking=no -o IdentitiesOnly=yes -i %s' % deploy_key_path
 
         with repo.git.custom_environment(GIT_SSH_COMMAND=git_ssh_cmd):
             local_branch = self.__base_name_local
-            if branch:
-                if self.__base_name_remote in branch:
-                    local_branch = branch.replace(f"{self.__base_name_remote}/", '').strip()
-                else:
-                    local_branch = branch
+            if repo.active_branch.name:
+                local_branch = repo.active_branch.name
 
-            repo.git.pull(self.__base_name_remote, local_branch)
+            origin = repo.remote(self.__base_name_remote)
+
+            if origin:
+                progress = ProgressPrinter()
+                try:
+                    origin.pull(refspec=local_branch, progress=progress)
+                    typer.echo(f"Pulled remote changes to branch '{local_branch}'")
+                except GitCommandError as ex:
+                    for line in progress.allDroppedLines():
+                        # returning the whole output if failed - so that user have any idea what's going on
+                        print(f'>> {line}')
+                    raise ex
 
     def find_main_branch_name(self, path: str) -> Optional[str]:
         repo = self.__get_repo(path=path)
@@ -275,10 +282,16 @@ class GitLocalClient:
         if repo:
             return repo.head.is_detached
 
-    def reset_hard(self, path: str):
+    def reset_hard(self, path: str, deploy_key_path: str, reset_to_origin: bool):
         repo = self.__get_repo(path=path)
         if repo:
-            repo.git.reset('--hard')
+            git_ssh_cmd = 'ssh -F /dev/null -o StrictHostKeyChecking=no -o IdentitiesOnly=yes -i %s' % deploy_key_path
+            with repo.git.custom_environment(GIT_SSH_COMMAND=git_ssh_cmd):
+                if reset_to_origin:
+                    repo.git.reset('--hard', f'origin/{repo.active_branch.name}')
+                    typer.echo(f'Branch "{repo.active_branch.name}" is now synced to its remote counterpart')
+                else:
+                    typer.echo('simple branch reset is not implemented')
 
     # refers to a "headless commit" where something was committed while in detached head state and head is pointing at that commit
     def is_head_committed_in_headless_state(self, path: str) -> bool:
@@ -297,6 +310,11 @@ class GitLocalClient:
             for branch in repo.heads:
                 if branch.name == branch_name:
                     return True
+
+            for ref in repo.remotes.origin.refs:
+                if ref.remote_head == branch_name:
+                    typer.echo(f'Found remote branch "{branch_name}"')
+                    return True
         return False
 
     def checkout_to_new_branch(self, path: str, branch_name: str):

{thestage-0.5.43 → thestage-0.5.44}/thestage/services/clients/thestage_api/api_client.py

@@ -361,7 +361,7 @@ class TheStageApiClient(TheStageApiClientCore):
         result = UserProfileResponse.model_validate(response) if response else None
         return result if result else None
 
-    def stop_task_on_instance(
+    def cancel_task(
             self,
             token: str,
             task_id: int,

{thestage-0.5.43 → thestage-0.5.44}/thestage/services/filesystem_service.py

@@ -83,6 +83,7 @@ class FileSystemServiceCore:
                 file.write(new_line)
                 file.write('\n')
 
+    # TODO remove this fucking useless shit
     def check_if_path_exist(self, file: str) -> bool:
         path = self.get_path(file, auto_create=False)
         if path.exists():

{thestage-0.5.43 → thestage-0.5.44}/thestage/services/logging/logging_service.py

@@ -151,7 +151,7 @@ class LoggingService(AbstractService):
             print_logs_task.cancel()
             if not input_task.result():  # result is only expected if ctrl+D triggered EOFError
                 typer.echo(f"\rTask {task_id} will be canceled")
-                self.__thestage_api_client.stop_task_on_instance(
+                self.__thestage_api_client.cancel_task(
                     token=config.main.thestage_auth_token,
                     task_id=task.id,
                 )

{thestage-0.5.43 → thestage-0.5.44}/thestage/services/project/project_service.py

@@ -148,11 +148,6 @@ class ProjectService(AbstractService):
 
         self.__git_local_client.git_fetch(path=config.runtime.working_directory, deploy_key_path=deploy_key_path)
 
-        branch = self.__git_local_client.find_main_branch_name(path=config.runtime.working_directory, )
-        if branch:
-            self.__git_local_client.git_pull(path=config.runtime.working_directory, deploy_key_path=deploy_key_path,
-                                             branch=branch)
-
         self.__git_local_client.init_gitignore(path=config.runtime.working_directory)
 
         self.__git_local_client.git_add_all(repo_path=config.runtime.working_directory)
@@ -401,6 +396,19 @@ class ProjectService(AbstractService):
             typer.echo("The task failed with an error")
             raise typer.Exit(1)
 
+    @error_handler()
+    def cancel_task(self, task_id: int, config: ConfigEntity):
+        cancel_result = self.__thestage_api_client.cancel_task(
+            token=config.main.thestage_auth_token,
+            task_id=task_id,
+        )
+
+        if cancel_result.is_success:
+            typer.echo(f'Task {task_id} has been canceled')
+        else:
+            typer.echo(f'Task {task_id} could not be canceled: {cancel_result.message}')
+
+
     @error_handler()
     def project_run_inference_simulator(
             self,
@@ -828,6 +836,8 @@ class ProjectService(AbstractService):
                 typer.echo("No main branch found")
 
 
+
+
     @error_handler()
     def set_default_container(self, config: ConfigEntity, container_uid: Optional[str]):
         project_config: ProjectConfig = self.__config_provider.read_project_config()
@@ -880,6 +890,8 @@ class ProjectService(AbstractService):
             typer.echo(f"No project found in working directory")
             raise typer.Exit(1)
 
+        is_deploy_key_exists = project_config.deploy_key_path and self.__file_system_service.check_if_path_exist(project_config.deploy_key_path)
+
         typer.echo(tabulate(
             [
                 [
@@ -888,11 +900,20 @@ class ProjectService(AbstractService):
                 [
                     "Default docker container unique ID", project_config.default_container_uid if project_config.default_container_uid else "<None>"
                 ],
+                [
+                    "Deploy key path", project_config.deploy_key_path if is_deploy_key_exists else "<None>"
+                ],
             ],
             showindex=False,
             tablefmt="simple",
         ))
 
+        if is_deploy_key_exists:
+            typer.echo("")
+            typer.echo(f"You can insert the following text:")
+            print(f"[{ColorScheme.USEFUL_INFO.value}]GIT_SSH_COMMAND=\"ssh -o StrictHostKeyChecking=no -o IdentitiesOnly=yes -i {project_config.deploy_key_path}\"[{ColorScheme.USEFUL_INFO.value}]")
+            typer.echo(f"before any regular git command to manage your local Project repository directly")
+
     @error_handler()
     def __get_fixed_project_config(self, config: ConfigEntity) -> Optional[ProjectConfig]:
         project_config: ProjectConfig = self.__config_provider.read_project_config()
@@ -1083,4 +1104,38 @@ class ProjectService(AbstractService):
 
         except Exception as e:
             typer.echo(__("Failed to deploy the inference simulator model to SageMaker: %error%", {"error": str(e)}))
-            raise typer.Exit(1)
+            raise typer.Exit(1)
+
+
+    @error_handler()
+    def pull_project(self, config: ConfigEntity):
+        project_config: ProjectConfig = self.__get_fixed_project_config(config=config)
+        if not project_config:
+            typer.echo(__("No project found at the path: %path%. Please initialize or clone a project first.", {"path": config.runtime.working_directory}))
+            raise typer.Exit(1)
+
+        typer.echo("Pulling code from remote repository...")
+        self.__git_local_client.git_pull(
+            path=config.runtime.working_directory,
+            deploy_key_path=project_config.deploy_key_path,
+        )
+
+
+    @error_handler()
+    def reset_project(self, config: ConfigEntity):
+        project_config: ProjectConfig = self.__get_fixed_project_config(config=config)
+        if not project_config:
+            typer.echo(__("No project found at the path: %path%. Please initialize or clone a project first.", {"path": config.runtime.working_directory}))
+            raise typer.Exit(1)
+
+        typer.echo("Fetching code from remote repository...")
+        self.__git_local_client.git_fetch(
+            path=config.runtime.working_directory,
+            deploy_key_path=project_config.deploy_key_path,
+        )
+        typer.echo("Resetting local branch...")
+        self.__git_local_client.reset_hard(
+            path=config.runtime.working_directory,
+            deploy_key_path=project_config.deploy_key_path,
+            reset_to_origin=True
+        )