TheStage AI Platform: CLI

Introduction

TheStage AI is an integrated platform designed for AI researchers, focusing on solving challenges related to accelerating deep neural network (DNN) inference.

TheStage CLI provides a command-line interface for interacting with TheStage AI infrastructure, allowing researchers to:

  • Server Instances: Retrieve information and connect to instances.

  • Docker Containers: Manage containers by retrieving information, connecting, streaming real-time logs, uploading/downloading files, and starting/stopping containers.

  • Projects: Clone and initialize projects, retrieve project specific task information, run tasks within a project, and stream real-time task logs.

TheStage AI’s Web interface provides full control over managing resources such as creating, renting, deleting components, etc..

Requirements

TheStage CLI is developed using Python and Poetry. Supported versions are Python 3.9 to 3.12 and Poetry 1.7.1. While the CLI is compatible with various operating systems, we recommend using a Unix-like OS such as Linux or macOS for optimal performance. For Windows users, consider running it through Windows Subsystem for Linux (WSL).

Additional Software Requirements:

  • Git

  • Python 3.9 to 3.12

  • SSH client

These tools must be installed for the CLI to function correctly.

Authentication and Authorization

To securely connect to TheStage AI’s backend and access information about server instances, containers, and more, API tokens are required. These tokens serve as both authentication and authorization mechanisms and can be generated through TheStage AI’s Web interface.

Note: An account on TheStage AI is required to generate an API token. Please refer to TheStage AI Platform: SSH Keys and API Tokens for more information.

Installation and Initialization

To generate an API token, please login to your TheStage AI account, and navigate to the “Profile > API tokens” section. Please refer to TheStage AI Platform: SSH Keys and API Tokens for more information.

# To install TheStage CLI:
pip install thestage

# To upgrade TheStage CLI:
pip install thestage --upgrade

# To set or change an API token:
thestage config set --api-token <API_token>

# To get help:
thestage --help

Using the CLI

Unique IDs

When working with components of your computational cluster — such as server instances, Docker containers, and projects — the CLI often requires you to specify the unique ID of the component. This unique ID is assigned by a user when the component is created in TheStage AI Web interface and can be found in your TheStage AI account or by listing the components using TheStage CLI.

Connecting to Server Instances and Containers

When connecting to a rented server instance or a Docker container hosted on it, the CLI uses the SSH key assigned to that server instance, which is stored in TheStage AI platform. If the instance status is “online” but the connection cannot be established, verify that an SSH key is associated with the server instance. SSH keys can be managed at TheStage AI account > Profile > SSH Keys.

When connecting to a self-hosted instance or a Docker container running on it, the CLI requires the username to be specified because it does not have information on which user to use. Ensure that the specified user has SSH access to the server. The instance status must be “online” for a successful connection. For more information, see TheStage AI Platform: Self-hosted Instances

Commands

Server Instances

TheStage AI enables researchers to rent computational resources from cloud providers or connect their self-hosted servers to a computational cluster on TheStage AI. With TheStage CLI, researchers can retrieve information and connect to both rented and connected server instances.

TheStage AI Web interface provides full control over managing server instances.

Rented Server Instances

Getting Information

# To get help:
thestage instance rented --help

# To list all rented server instances:
thestage instance rented ls -s all

# To filter rented server instances by status, for example 'online':
thestage instance rented ls -s online

# To see other options:
thestage instance rented ls --help

Connecting to a rented server Instance

# There are two ways to connect to a rented server instance:
thestage instance rented connect <instance_unique_ID>
thestage connect <instance_unique_ID>

Self-hosted Instances

Getting Information

# To get help:
thestage instance self-hosted --help

# To list all self-hosted instances:
thestage instance self-hosted ls -s all

# To filter self-hosted instances by status, for example 'awaiting_setup':
thestage instance self-hosted ls -s awaiting_setup

# To see other options:
thestage instance self-hosted ls --help

Connecting to a self-hosted Instance

# There are two ways to connect to a self-hosted instance:
thestage instance self-hosted connect <instance_unique_ID> --username <username>
thestage connect <instance_unique_ID> --username <username>

Docker Containers

TheStage AI provides an interface for managing Docker containers within a computational cluster consisting of both rented and self-hosted servers. While complete container management can be performed using TheStage AI Web interface, TheStage CLI allows the following:

  • Retrieve information about Docker containers.

  • Manage containers, including starting and stopping them.

  • Connect to a Docker container.

  • Upload files to a Docker container.

  • Download files form a Docker container.

  • Stream real-time container logs.

Getting Information

# To get help:
thestage container --help

# To list all containers:
thestage container ls -s all

# To filter containers by status, for example 'running':
thestage container ls -s running

# To get additional information on a container:
thestage container info <container_unique_ID>

Starting/Stopping a container

# To start a container:
thestage container start <container_unique_ID>

# To stop a container:
thestage container stop <container_unique_ID>

Connecting to a container

When connecting to containers via SSH, the CLI first connects to the server instance hosting the container and then connects to the container.

# To connect to a container hosted on a rented server instance:
thestage container connect <container_unique_ID>
thestage connect <container_unique_ID>

# To connect to a container hosted on a self-hosted instance:
thestage container connect <container_unique_ID> --username <username>
thestage connect <container_unique_ID> --userneme <username>

Uploading files to a container

When uploading a file to a container, the CLI connects to the server instance hosting the container, uploads the file to a directory on the instance mapped to the container, and then moves the file to the specified destination directory inside the container.

# To upload a file to a container hosted on a rented server instance:
thestage container upload 'source_path_file' <container_unique_id>:'destination_dir_path'

# To upload a file to a container hosted on a self-hosted instance:
thestage container upload 'source_path_file' <container_unique_id>:'destination_dir_path' --username <username>

Downloads files from a container

When downloading a file from a container, the CLI connects to the server instance hosting the container, connects to the container, copies the specified file in the container to a directory mapped to the host, and then downloads the file to the specified directory on the local computer.

# to download a file from a container hosted on a rented server instance:
thestage container download <container_uid>:'source_path_file' 'destination_dir_path'

# to download a file from a container hosted on a self-hosted instance:
thestage container download <container_uid>:'source_path_file' 'destination_dir_path' --username <username>

Streaming container logs

TheStage CLI facilitates to stream real-time logs. Log history is available in your TheStage AI account.

# To stream real-time container logs:
thestage container logs <container_unique_ID>

Projects

Projects allow researchers to divide computational resources between teams or research directions and organize work within each project for better management and execution.

TheStage CLI provides the following functionalities:

  • Clone and initialize projects,

  • Retrieve project specific task information,

  • Run tasks within a project, and stream real-time task logs.

Project Workflow

Step 1: Setting up the Infrastructure (via TheStage AI Web Interface)

  1. Rent Servers or Connect Self-hosted Servers: Rent server instances from cloud providers or connect self-hosted servers to your TheStage AI account.

  2. Create a Project: Create a project on one or multiple servers. When a project is created, a dedicated GitHub repository is automatically generated for the project.

  3. Set Up Containers: Create containers specific to the project to run tasks within the project.

Step 2: Managing the Project with TheStage CLI

  1. Initialize or Clone the Project: Use the CLI to initialize (init) or clone (clone) the project repository locally. The thestage init and thestage clone commands are similar to the git init and git clone commands.

    • If you have a local directory with existing code and the project repository is empty, use the init command to start tracking your code in the project repository.

    • If the project repository already contains code, use the clone command to download it to your local directory.

  2. Create and Run Tasks: Create and run tasks within the project on the available containers. If there are any code changes, they will be committed to the project’s GitHub repository before task execution.

Real-Time Logs and Task Management

  • Real-Time Logs: Access real-time logs for tasks and containers through TheStage CLI.

  • Full Logs and History: Full task logs, run history, and commit history are available through the TheStage AI Web Interface.

  • Task Termination: Tasks can be terminated through the TheStage AI Web Interface. This functionality will soon be available in TheStage CLI as well.

Initializing/Cloning Project

# To initialize a project in a working directory::
thestage project init --working-directory 'dir_path' <project_unique_ID>
thestage project init -wd 'dir_path' <project_unique_ID>

# To clone a project to a working directory
thestage project clone --working-directory 'dir_path' <project_unique_ID>
thestage project clone -wd 'dir_path' <project_unique_ID>

Working with Tasks within a Project

# To create, submit changes, and run task within a project:
thestage project run --commit-hash 'commit_hash' --task-tile 'task_title' --container-uid <contaner_unique_ID> --working-directory 'project_dir_path' 'command_to_run'
thestage project run -hash 'commit_hash' -t 'task_title' -cid <contaner_unique_ID> -wd 'project_dir_path' 'command_to_run'

Parameters:

  • --commit-hash or -hash (optional): Specifies the commit hash to use for the task. If not provided, the latest commit from the main branch (HEAD commit) will be used.

  • --task-title or -t (optional): Assigns a title to the task for easier identification and tracking. If not provided, the task title will be empty.

Task Execution and Log Streaming:

  • When a task is initiated, the CLI by default streams the real-time logs for that task.

  • To terminate the task, press Ctrl+C. This will immediately stop the task execution. A running task can also be stopped using TheStage AI Web Interface.

  • To detach from the log stream without stopping the task, press Ctrl+D. The task will continue running in the background. To reattach to the log stream, run the following command: thestage project task logs <task_ID>

  • Full task logs, run history, and commit history are available through TheStage AI Web Interface.

Task Queuing:

  • If multiple tasks are queued to run in the same container, they will execute sequentially. A new task will not start until the previous one has completed. TheStage daemon running on the server instance manages the task queue.