first add v3.19.0

.coveragerc

@@ -0,0 +1,28 @@
+[run]
+omit = tests/*, aim/storage/migrations/*, aim/storage/env.py
+
+[report]
+# Regexes for lines to exclude from consideration
+exclude_lines =
+    # Have to re-enable the standard pragma
+    pragma: no cover
+
+    # Don't complain about missing debug-only code:
+    def __repr__
+    if self\.debug
+
+    # Don't complain if tests don't hit defensive assertion code:
+    raise AssertionError
+    raise NotImplementedError
+
+    # Don't complain if non-runnable code isn't run:
+    if 0:
+    if __name__ == .__main__.:
+
+    # Don't complain about abstract methods, they aren't run:
+    @(abc\.)?abstractmethod
+
+    if TYPE_CHECKING:
+
+[html]
+directory = coverage_html_report

.gitattributes

@@ -0,0 +1,3 @@
+*.ts linguist-detectable=false
+*.tsx linguist-detectable=false
+*.scss linguist-detectable=false

.gitignore

@@ -0,0 +1,125 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+*.c
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+pip-wheel-metadata/
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+docs/build/
+docs/source/generated
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+aim.sublime-*
+.aim/
+.aim-test/
+
+.DS_Store
+.idea
+.vim
+
+data
+examples/.aim
+examples/.git
+
+coverage_html_report/
+
+# generated cpp files
+# @TODO remove this line once there will be usage of cpp files
+aim/storage/**/*.cpp

.readthedocs.yaml

@@ -0,0 +1,23 @@
+# .readthedocs.yaml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+# Build documentation in the docs/ directory with Sphinx
+# reference: https://docs.readthedocs.io/en/stable/config-file/v2.html#sphinx
+sphinx:
+  configuration: docs/source/conf.py
+  fail_on_warning: false
+
+# Optionally build your docs in additional formats such as PDF
+formats:
+   - pdf
+   
+# Optionally set the version of Python and requirements required to build your docs
+python:
+  version: 3.7
+  install:
+    - requirements: requirements.dev.txt
+    - requirements: docs/requirements.txt

CITATION.cff

@@ -0,0 +1,16 @@
+cff-version: 1.2.0
+message: "If you want to cite the tool, feel free to use this (but only if you loved it 😊)"
+title: "Aim"
+abstract: "An easy-to-use and performant open-source ML experiment tracking tool."
+date-released: 2020-06-18
+authors:
+  - family-names: "Arakelyan"
+    given-names: "Gor"
+  - family-names: "Soghomonyan"
+    given-names: "Gevorg"
+  - name: "The Aim team"
+version: 3.9.3
+doi: 10.5281/zenodo.6536395
+license: "Apache-2.0"
+url: "https://aimstack.io"
+repository-code: "https://github.com/aimhubio/aim"

CODE_OF_CONDUCT.md

@@ -0,0 +1,74 @@
+## Code of Conduct
+
+### Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to making participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, gender identity and expression, level of experience,
+nationality, personal appearance, race, religion, or sexual identity and
+orientation.
+
+### Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+  address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+### Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+### Scope
+
+This Code of Conduct applies both within project spaces and in public spaces
+when an individual is representing the project or its community. Examples of
+representing a project or community include using an official project e-mail
+address, posting via an official social media account, or acting as an appointed
+representative at an online or offline event. Representation of a project may be
+further defined and clarified by project maintainers.
+
+### Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at community@aimstack.io. All
+complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+### Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
+available at [http://contributor-covenant.org/version/1/4][version]
+
+[homepage]: http://contributor-covenant.org
+[version]: http://contributor-covenant.org/version/1/4/

CONTRIBUTING.md

@@ -0,0 +1,153 @@
+# Contributing
+
+Thanks for your interest in helping improve Aim! 🎉
+
+## Before Contributing
+
+As with most projects, prior to starting to code on a bug fix or feature request, please post in the respective GitHub issue saying you want to volunteer, and then wait for a positive response. And if there is no issue for it yet, create it first.
+
+This helps make sure:
+1. Two people aren't working on the same thing.
+2. This is something Aim's maintainers believe should be implemented/fixed.
+3. Any API, UI, or deeper architectural changes that need to be implemented have been fully thought through by the community together with Aim's maintainers.
+
+Please follow [Aim Contributor Covenant Code of Conduct](./CODE_OF_CONDUCT.md) in all your interactions with the project.
+
+## Governance
+
+This section describes governance processes we follow in developing Aim.
+
+### Persons of Interest
+
+#### Authors
+
+- Gev Soghomonyan ([SGevorg](https://github.com/SGevorg))
+- Gor Arakelyan ([gorarakelyan](https://github.com/gorarakelyan))
+
+#### Primary Maintainers
+
+- Albert Torosyan ([alberttorosyan](https://github.com/alberttorosyan))
+- Karen Hambardzumyan ([mahnerak](https://github.com/mahnerak))
+- Ruben Karapetyan ([roubkar](https://github.com/roubkar))
+
+#### Core Contributors
+
+- Arsen Alkamyan ([arsengit](https://github.com/arsengit))
+- Karo Muradyan ([KaroMourad](https://github.com/KaroMourad))
+- Mihran Vanyan ([mihran113](https://github.com/mihran113))
+
+### Releases
+
+We release a new minor version (e.g., 3.8.0) every three to four week and patch releases on demand. The minor versions contain new features, bugfixes and also all previous bugfixes included in previous patch releases. With every release, we publish a [CHANGELOG](./CHANGELOG.md) where we list enhancements and fixes. The versioning scheme we use is [SemVer](http://semver.org/).
+
+## Contribution Process
+
+The Aim contribution process starts with filing a GitHub issue. Aim defines six categories of issues: enhancements (feature requests), bug reports, code health improvements, peformance improvements, tests, questions.
+
+Aim maintainers actively triage and respond to GitHub issues. In general, we recommend waiting for feedback from an Aim maintainer or community member before proceeding to implement a feature or patch. This is particularly important for significant changes, and will typically be labeled during triage with `phase / exploring`.
+
+After you have agreed upon an implementation strategy for your feature or patch with an Aim maintainer, the next step is to introduce your changes as a pull request against the Aim Repository.
+
+Once your pull request against the Aim Repository has been merged, your corresponding changes will be automatically included in the next Aim release. Every change is listed in the [Aim release notes](https://github.com/aimhubio/aim/releases) and [CHANGELOG](./CHANGELOG.md).
+
+Congratulations, you have just contributed to Aim. We appreciate your contribution!
+
+## Developing and Testing
+
+The majority of the Aim product areas is developed in Python/Cython. This includes the Storage, SDK, Tracking Server, CLI, API. Aim UI is a Web app mostly built with TypeScript and React.
+
+### Developing Storage/SDK/CLI
+
+Most of the backend components, including SDK, Storage, Web APIs and CLI are developed using Python/Cython.
+In order to start development you must install dev requirements and the aim package itself, in editable mode.
+```shell
+pip install -r requirements.dev.txt
+pip install -e .
+```
+
+Verify that Aim installed properly by running
+```shell
+aim version
+```
+or by importing aim in python REPL
+```python
+import aim
+```
+
+#### Style Guide
+Aim follows PEP8 standard for style guide and uses `flake8` as a style checker. Style checks enforced
+as a check on GitHub Actions when new PR opened.
+
+#### Testing Requirements
+
+New unit-tests must be added along with the code changes. In order to setup the testing environment
+```shell
+cd tests
+pip install -r requirements.txt
+```
+
+Aim python code unit-tests are located at `tests/` directory. Unit-tests are written in Python's `unittest` package style.
+[Pytest](https://docs.pytest.org) is used as a test runner/discovery tool. To make sure unit-tests are not failing run
+```shell
+pytest tests/
+```
+
+### Developing UI
+
+Aim UI is written in TypeScript. `npm` is required to build Aim UI and to run in DEV mode.
+You can verify that `npm` is on the PATH by running `npm -v`, and
+[install npm](https://www.npmjs.com/get-npm) if needed.
+
+#### Style Guide
+
+We use Prettier to autoformat code on presubmit.
+
+#### Launching the Development UI
+
+Before running the Aim UI dev server or building a distributable wheel, install npm
+dependencies via:
+
+```shell
+cd aim/web/ui
+npm install
+```
+
+Then you can start the dev server:
+
+```shell
+npm start
+```
+
+Aim UI will show logged data in at [http://localhost:3000](http://localhost:3000).
+
+#### Adding New Components
+
+To start building a new component you can run following command:
+
+```shell
+npm run crc 'ComponentName'
+```
+
+
+If you want to add a component inside UI kit you can run following command:
+
+```shell
+npm run crc-kit 'ComponentName'
+```
+
+These command will create a folder named `ComponentName` with all the necessary files.
+
+## Writing Docs
+
+Aim documentation is built using [Sphix](https://www.sphinx-doc.org) and is hosted at
+[Read the Docs](https://aimstack.readthedocs.io).
+The documentation sources are located at `docs/` directory. In order to build documentation locally
+run the following commands
+```shell
+pip install -r requirements.dev.txt
+cd docs
+pip install -r requirements.txt
+make html
+```
+
+Documentation will be available at `docs/build/html/index.html` on your local machine.

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2019-2022 Aimhub, Inc.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

MANIFEST.in

@@ -0,0 +1,5 @@
+include README.md LICENSE aim/VERSION
+recursive-include aim/storage *.h *.hpp *.cpp *.cxx *.pxd *.pyx
+recursive-include aim/storage/migrations *
+recursive-include aim/web/migrations *
+

aim/VERSION

@@ -0,0 +1 @@
+3.19.0

aim/__about__.py

@@ -0,0 +1,26 @@
+import time
+import os
+
+here = os.path.abspath(os.path.dirname(__file__))
+
+version_file = f'{here}/VERSION'
+
+__version__ = None
+with open(version_file) as vf:
+    __version__ = vf.read().strip()
+
+_this_year = time.strftime("%Y")
+__author__ = "Gev Sogomonian, Gor Arakelyan et al."
+__author_email__ = "community@aimhub.io"
+__license__ = "Apache-2.0"
+__copyright__ = f"2019-{_this_year}, {__author__}."
+__homepage__ = "https://github.com/aimhubio/aim"
+__docs_url__ = "https://github.com/aimhubio/aim"
+# this has to be simple string, see: https://github.com/pypa/twine/issues/522
+__docs__ = (
+    "Aim is a lightweight library for ML researchers to store and retrieve AI metadata."
+)
+__long_docs__ = """
+"""
+
+__all__ = ["__author__", "__author_email__", "__copyright__", "__docs__", "__homepage__", "__license__", "__version__"]

aim/__init__.py

@@ -0,0 +1,8 @@
+import aimrocks
+from aim.sdk import *
+from aim.ext.notebook.notebook import load_ipython_extension
+from aim.cli.manager.manager import run_process
+
+from aim.utils.deprecation import python_version_deprecation_check
+
+python_version_deprecation_check()

aim/__version__.py

@@ -0,0 +1,10 @@
+import os
+
+here = os.path.abspath(os.path.dirname(__file__))
+
+version_file = f'{here}/VERSION'
+
+
+__version__ = None
+with open(version_file) as vf:
+    __version__ = vf.read().strip()

aim/acme.py

@@ -0,0 +1,2 @@
+# Alias to SDK acme interface
+from aim.sdk.adapters.acme import AimCallback, AimWriter  # noqa F401

aim/catboost.py

@@ -0,0 +1,2 @@
+# Alias to SDK Catboost interface
+from aim.sdk.adapters.catboost import AimLogger  # noqa: F401

aim/cli/cli.py

@@ -0,0 +1,31 @@
+import click
+from click import core
+
+from aim.cli.configs import *  # noqa F403
+from aim.cli.init import commands as init_commands
+from aim.cli.version import commands as version_commands
+from aim.cli.up import commands as up_commands
+from aim.cli.server import commands as server_commands
+from aim.cli.reindex import commands as reindex_commands
+from aim.cli.runs import commands as runs_commands
+from aim.cli.convert import commands as convert_commands
+from aim.cli.storage import commands as storage_commands
+
+core._verify_python3_env = lambda: None
+
+
+@click.group()
+@click.option('-v', '--verbose', is_flag=True)
+def cli_entry_point(verbose):
+    if verbose:
+        click.echo('Verbose mode is on')
+
+
+cli_entry_point.add_command(init_commands.init, INIT_NAME)
+cli_entry_point.add_command(version_commands.version, VERSION_NAME)
+cli_entry_point.add_command(up_commands.up, UP_NAME)
+cli_entry_point.add_command(server_commands.server, SERVER_NAME)
+cli_entry_point.add_command(reindex_commands.reindex, REINDEX_NAME)
+cli_entry_point.add_command(runs_commands.runs, RUNS_NAME)
+cli_entry_point.add_command(convert_commands.convert, CONVERT)
+cli_entry_point.add_command(storage_commands.storage, STORAGE)

aim/cli/configs.py

@@ -0,0 +1,8 @@
+INIT_NAME = 'init'
+VERSION_NAME = 'version'
+UP_NAME = 'up'
+REINDEX_NAME = 'reindex'
+SERVER_NAME = 'server'
+RUNS_NAME = 'runs'
+CONVERT = 'convert'
+STORAGE = 'storage'

aim/cli/convert/commands.py

@@ -0,0 +1,76 @@
+import os
+
+import click
+from click import ClickException
+
+from aim.sdk.repo import Repo
+from aim.sdk.utils import clean_repo_path
+from aim.cli.convert.processors import (
+    parse_tb_logs,
+    parse_mlflow_logs,
+    parse_wandb_logs,
+)
+
+
+@click.group()
+@click.option('--repo', required=False, type=click.Path(exists=True,
+                                                        file_okay=False,
+                                                        dir_okay=True,
+                                                        writable=True))
+@click.pass_context
+def convert(ctx, repo):
+    ctx.ensure_object(dict)
+
+    repo_path = clean_repo_path(repo) or Repo.default_repo_path()
+    repo_inst = Repo.from_path(repo_path)
+
+    ctx.obj['repo_inst'] = repo_inst
+
+
+@convert.command(name='tensorboard')
+@click.pass_context
+@click.option('--logdir', required=True, type=click.Path(exists=True,
+                                                         readable=True,
+                                                         dir_okay=True,
+                                                         resolve_path=True))
+@click.option('--flat', '-f', required=False, is_flag=True, default=False)
+@click.option('--no-cache', required=False, is_flag=True, default=False)
+def convert_tensorboard(ctx, logdir, flat, no_cache):
+    repo_inst = ctx.obj['repo_inst']
+    parse_tb_logs(logdir, repo_inst, flat, no_cache)
+
+
+@convert.command(name='tf')
+@click.pass_context
+@click.option('--logdir', required=True, type=click.Path(exists=True,
+                                                         readable=True,
+                                                         dir_okay=True,
+                                                         resolve_path=True))
+@click.option('--flat', '-f', required=False, is_flag=True, default=False)
+def convert_tensorflow(ctx, logdir, flat):
+    click.secho('WARN: Command \'tf\' is deprecated and will be removed in future releases,'
+                ' please use \'tensorboard\' instead.', fg='red')
+    repo_inst = ctx.obj['repo_inst']
+    parse_tb_logs(logdir, repo_inst, flat)
+
+
+@convert.command(name='mlflow')
+@click.pass_context
+@click.option('--tracking_uri', required=False, default=None)
+@click.option('--experiment', '-e', required=False, default=None)
+def convert_mlflow(ctx, tracking_uri=None, **kwargs):
+    repo_inst = ctx.obj['repo_inst']
+    tracking_uri = tracking_uri or os.environ.get("MLFLOW_TRACKING_URI")
+    if not tracking_uri:
+        raise ClickException("MLFlow tracking URI must be provided either trough ENV or CLI.")
+    parse_mlflow_logs(repo_inst, tracking_uri, **kwargs)
+
+
+@convert.command(name='wandb')
+@click.pass_context
+@click.option('--entity', required=True, default=None)
+@click.option('--project', required=True, default=None)
+@click.option('--run-id', required=False, default=None)
+def convert_wandb(ctx, entity=None, project=None, **kwargs):
+    repo_inst = ctx.obj['repo_inst']
+    parse_wandb_logs(repo_inst, entity, project, **kwargs)

aim/cli/convert/processors/__init__.py

@@ -0,0 +1,3 @@
+from .tensorboard import parse_tb_logs
+from .mlflow import parse_mlflow_logs
+from .wandb import parse_wandb_logs

aim/cli/convert/processors/mlflow.py

@@ -0,0 +1,152 @@
+import os.path
+from tempfile import TemporaryDirectory
+
+import click
+
+from aim import Run, Image, Text, Audio
+
+IMAGE_EXTENSIONS = ('jpg', 'bmp', 'jpeg', 'png', 'gif', 'svg')
+HTML_EXTENSIONS = ('html',)
+TEXT_EXTENSIONS = (
+    'txt',
+    'log',
+    'py',
+    'js',
+    'yaml',
+    'yml',
+    'json',
+    'csv',
+    'tsv',
+    'md',
+    'rst',
+    'jsonnet',
+)
+
+# Audio is not handled in mlflow but including here just in case
+AUDIO_EXTENSIONS = (
+    'flac',
+    'mp3',
+    'wav',
+)
+
+
+def parse_mlflow_logs(repo_inst, tracking_uri, experiment):
+    try:
+        import mlflow
+    except ImportError:
+        click.echo(
+            'Could not process mlflow logs - failed to import "mlflow" module.', err=True
+        )
+        return
+
+    client = mlflow.tracking.client.MlflowClient(tracking_uri=tracking_uri)
+
+    if experiment is None:
+        # process all experiments
+        experiments = client.list_experiments()
+    else:
+        try:
+            ex = client.get_experiment(experiment)
+        except mlflow.exceptions.MlflowException:
+            ex = client.get_experiment_by_name(experiment)
+        if not ex:
+            click.echo(f'Could not find experiment with id or name "{experiment}"', err=True)
+            return
+        experiments = (ex,)
+
+    for ex in experiments:
+        runs = client.search_runs(ex.experiment_id)
+        for run in runs:
+            run_id = run.info.run_id
+            aim_run = Run(
+                repo=repo_inst,
+                system_tracking_interval=None,
+                capture_terminal_logs=False,
+                experiment=ex.experiment_id,
+            )
+            aim_run['mlflow_run_id'] = run.info.run_id
+            aim_run['mlflow_run_name'] = run.data.tags.get("mlflow.runName")
+            aim_run.description = run.data.tags.get("mlflow.note.content")
+
+            # Collect params & tags
+            aim_run['params'] = run.data.params
+            aim_run['tags'] = {
+                k: v for k, v in run.data.tags.items() if not k.startswith('mlflow')
+            }
+
+            # Collect metrics
+            for key in run.data.metrics.keys():
+                for m in client.get_metric_history(run_id, key):
+                    aim_run.track(m.value, step=m.step, name=m.key)
+
+            # Collect artifacts
+            __html_warning_issued = False
+            with TemporaryDirectory(prefix=f'mlflow_{run.info.run_id}_') as temp_path:
+                # click.secho(f'Downloading artifacts to {temp_path}', fg='green')
+                artifact_loc_stack = [None]
+                while artifact_loc_stack:
+                    loc = artifact_loc_stack.pop()
+                    artifacts = client.list_artifacts(run_id, path=loc)
+
+                    img_batch = []
+                    text_batch = []
+                    audio_batch = []
+
+                    for file_info in artifacts:
+                        if file_info.is_dir:
+                            artifact_loc_stack.append(file_info.path)
+                            continue
+
+                        downloaded_path = client.download_artifacts(run_id, file_info.path, dst_path=temp_path)
+                        if file_info.path.endswith(HTML_EXTENSIONS):
+                            # TODO [AP] plotly does not provide interface to load from html
+                            # TODO [AP] need to implement html custom object?
+                            if not __html_warning_issued:
+                                click.secho(
+                                    'Handler for html file types is not yet implemented.', fg='yellow'
+                                )
+                                __html_warning_issued = True
+                            continue
+                        elif file_info.path.endswith(IMAGE_EXTENSIONS):
+                            aim_object = Image
+                            kwargs = dict(
+                                image=downloaded_path,
+                                caption=file_info.path
+                            )
+                            container = img_batch
+                        elif file_info.path.endswith(TEXT_EXTENSIONS):
+                            with open(downloaded_path) as fh:
+                                content = fh.read()
+                            aim_object = Text
+                            kwargs = dict(
+                                text=content
+                            )
+                            container = text_batch
+                        elif file_info.path.endswith(AUDIO_EXTENSIONS):
+                            audio_format = os.path.splitext(file_info.path)[1].lstrip('.')
+                            aim_object = Audio
+                            kwargs = dict(
+                                data=downloaded_path,
+                                caption=file_info.path,
+                                format=audio_format
+                            )
+                            container = audio_batch
+                        else:
+                            click.secho(
+                                f'Unresolved or unsupported type for artifact {file_info.path}', fg='yellow'
+                            )
+                            continue
+
+                        try:
+                            item = aim_object(**kwargs)
+                        except Exception as exc:
+                            click.echo(
+                                f'Could not convert artifact {file_info.path} into aim object - {exc}', err=True
+                            )
+                            continue
+                        container.append(item)
+
+                    for content_type, seq in (('image', img_batch),
+                                              ('text', text_batch),
+                                              ('audio', audio_batch)):
+                        aim_run.track(seq, step=0, name=loc or 'root', context={'type': content_type})

aim/cli/convert/processors/tensorboard.py

@@ -0,0 +1,262 @@
+import json
+import os
+
+import click
+from tqdm import tqdm
+
+from aim import Audio, Image, Run
+
+
+def parse_tb_logs(tb_logs, repo_inst, flat=False, no_cache=False):
+    """
+    This function scans and collects records from TB log files.
+
+    Creates and uses cache file "tb_logs_cache" in the repo dir
+    to track previously processed files and values
+
+    For more info please refer to our integration guides.
+    """
+
+    try:
+        # This import statement takes long to complete
+        import tensorflow as tf
+        from tensorflow.python.summary.summary_iterator import summary_iterator
+        from tensorboard.util import tensor_util
+    except ImportError:
+        click.echo(
+            'Could not process TensorBoard logs - failed to import tensorflow module.', err=True
+        )
+        return
+
+    supported_plugins = ('images', 'scalars')
+    unsupported_plugin_noticed = False
+    tb_logs_cache_path = os.path.join(repo_inst.path, 'tb_logs_cache')
+
+    if no_cache and os.path.exists(tb_logs_cache_path):
+        os.remove(tb_logs_cache_path)
+    try:
+        with open(tb_logs_cache_path) as FS:
+            tb_logs_cache = json.load(FS)
+    except Exception:
+        tb_logs_cache = {}
+
+    def get_parent(current_path, level=0):
+        # level 0 is the direct parent directory
+        if level <= 0:
+            return os.path.dirname(current_path)
+        elif current_path in ('', '.', '/'):
+            return current_path
+        return get_parent(os.path.dirname(current_path), level - 1)
+
+    tb_logs = os.path.abspath(tb_logs)
+    run_dir_candidates = set()
+    for root, dirs, files in os.walk(tb_logs):
+        for file in files:
+            if not file.startswith('events.out.tfevents'):
+                continue
+
+            file_path = os.path.abspath(os.path.join(root, file))
+            run_dir = get_parent(file_path)
+
+            if not run_dir.startswith(tb_logs):
+                # it's outside tb_logs
+                continue
+
+            run_dir_candidates.add(run_dir)
+
+    def get_level(current_path):
+        level = -1
+        while current_path.startswith(tb_logs):
+            current_path, _ = os.path.split(current_path)
+            level += 1
+        return level
+
+    def create_ndarray(tensor):
+        res = tensor_util.make_ndarray(tensor)
+        if res.dtype == "object":
+            return None
+        else:
+            return res
+
+    run_dir_candidates = sorted(run_dir_candidates, key=get_level, reverse=True)
+    run_dir_candidates_filtered = set()
+    run_dir_ignored = set()
+    groups = set()
+
+    for run_dir in run_dir_candidates:
+        if run_dir in run_dir_candidates_filtered:
+            # already tagged as a run dir
+            continue
+
+        if run_dir in groups:
+            # run dir which has other run dirs inside, so we skip it
+            run_dir_ignored.add(run_dir)
+            continue
+
+        depth = get_level(run_dir)
+        if depth >= 2:
+            if flat:
+                run_group_dir = get_parent(run_dir, 0)
+                new_run_dir = run_dir
+            else:
+                run_group_dir = get_parent(run_dir, 1)
+                new_run_dir = get_parent(run_dir, 0)
+                if new_run_dir in groups:
+                    new_run_dir = run_dir
+            groups.add(run_group_dir)
+        elif depth == 1:
+            new_run_dir = run_dir
+        else:
+            continue
+        run_dir_candidates_filtered.add(new_run_dir)
+
+    if run_dir_ignored:
+        click.echo('WARN: Found directory entries with unorganized even files!\n'
+                   'Please read the preparation instructions to properly process these files.\n'
+                   'Event files in the following directories will be ignored:', err=True)
+        for c, r in enumerate(run_dir_ignored, start=1):
+            click.echo(f'{c}: {r}', err=True)
+
+    for path in tqdm(run_dir_candidates_filtered,
+                     desc='Converting TensorBoard logs',
+                     total=len(run_dir_candidates_filtered)):
+
+        events = {}
+        for root, dirs, files in os.walk(path):
+            for file in files:
+                if 'events.out.tfevents' not in file:
+                    continue
+                file_path = os.path.join(root, file)
+                if file_path == os.path.join(path, file):
+                    entry = None
+                else:
+                    entry = os.path.basename(os.path.dirname(file_path))
+                events[file_path] = {
+                    'context': {
+                        'entry': entry
+                    }
+                }
+
+        if path not in tb_logs_cache:
+            tb_logs_cache[path] = {}
+
+        run_cache = tb_logs_cache[path]
+        if run_cache:
+            run = Run(
+                run_hash=run_cache['run_hash'],
+                repo=repo_inst,
+                system_tracking_interval=None,
+                log_system_params=False,
+                capture_terminal_logs=False,
+            )
+        else:
+            run = Run(
+                repo=repo_inst,
+                system_tracking_interval=None,
+                log_system_params=False,
+                capture_terminal_logs=False,
+            )
+            run['tensorboard_logdir'] = path
+            run_cache.update({
+                'run_hash': run.hash,
+                'events': {},
+            })
+        run_tb_events = run_cache['events']
+
+        events_to_process = []
+        for event in events:
+            last_modified_at = os.path.getmtime(event)
+            try:
+                assert last_modified_at == run_tb_events[event]['last_modified_at']
+            except (KeyError, AssertionError, RuntimeError):
+                # Something has changed or hasn't been processed before
+                events_to_process.append(event)
+                try:
+                    run_tb_events[event]['last_modified_at'] = last_modified_at
+                except KeyError:
+                    # Completely new event
+                    run_tb_events[event] = {
+                        'last_modified_at': last_modified_at,
+                        'values': {},
+                    }
+
+        if not events_to_process:
+            continue
+
+        for event_file in tqdm(events_to_process, desc=f'Parsing logs in {path}', total=len(events_to_process)):
+            run_tb_log = run_tb_events[event_file]
+            event_context = events[event_file]['context']
+            try:
+                for event in summary_iterator(event_file):
+                    timestamp = event.wall_time
+                    step = event.step
+                    fail_count = 0
+                    _err_info = None
+
+                    for value in event.summary.value:
+                        tag = value.tag
+
+                        plugin_name = value.metadata.plugin_data.plugin_name
+                        value_id = f'{tag}_{plugin_name}'
+                        if value_id in run_tb_log['values']:
+                            if run_tb_log['values'][value_id]['timestamp'] >= timestamp:
+                                # prevent previously tracked data from re-tracking upon file update
+                                continue
+
+                        if len(plugin_name) > 0 and plugin_name not in supported_plugins:
+                            if not unsupported_plugin_noticed:
+                                click.echo(
+                                    'Found unsupported plugin type in the log file. '
+                                    'Data for these wont be processed. '
+                                    'Supported plugin types are: {}'.format(', '.join(supported_plugins)),
+                                    err=True
+                                )
+                                unsupported_plugin_noticed = True
+                            continue
+                        track_val = None
+                        try:
+                            if value.HasField('tensor'):
+                                # TODO: [MV] check the case when audios are passed via tensor
+                                if plugin_name == 'images':
+                                    tensor = value.tensor.string_val[2:]
+                                    track_val = [
+                                        Image(tf.image.decode_image(t).numpy()) for t in tensor
+                                    ]
+                                    if len(track_val) == 1:
+                                        track_val = track_val[0]
+                                elif plugin_name == "scalars" or plugin_name == "":
+                                    track_val = create_ndarray(value.tensor)
+                                else:
+                                    track_val = value.tensor.float_val[0]
+                            elif value.HasField('simple_value'):
+                                track_val = value.simple_value
+                            elif value.HasField('image'):
+                                track_val = Image(tf.image.decode_image(value.image.encoded_image_string).numpy())
+                            elif value.HasField('audio'):
+                                tf_audio, sample_rate = tf.audio.decode_wav(value.audio.encoded_audio_string)
+                                track_val = Audio(tf_audio.numpy(), rate=sample_rate)
+
+                        except RuntimeError as exc:
+                            # catch all the nasty failures
+                            fail_count += 1
+                            if not _err_info:
+                                _err_info = str(exc)
+                            continue
+
+                        run_tb_log['values'][value_id] = {
+                            'step': step,
+                            'timestamp': timestamp
+                        }
+                        if track_val is not None:
+                            run._tracker._track(track_val, timestamp, tag, step, context=event_context)
+                    if fail_count:
+                        click.echo(f'Failed to process {fail_count} entries. First exception: {_err_info}', err=True)
+
+            except RuntimeError as exc:
+                click.echo(f'Failed to read log file {event_file} - {exc}', err=True)
+
+    # refresh cache
+    with open(tb_logs_cache_path, 'w') as FS:
+        json.dump(tb_logs_cache, FS)
+
+    click.echo('TensorBoard logs conversion complete!')

aim/cli/convert/processors/wandb.py

@@ -0,0 +1,182 @@
+from pathlib import Path
+import re
+from tempfile import TemporaryDirectory
+
+import click
+from tqdm import tqdm
+
+from aim import Run
+from aim.ext.resource.log import LogLine
+from aim.ext.resource.configs import AIM_RESOURCE_METRIC_PREFIX
+
+
+def parse_wandb_logs(repo_inst, entity, project, run_id):
+    try:
+        import wandb
+    except ImportError:
+        click.echo("Could not process wandb logs - failed to import 'wandb' module.", err=True)
+        return
+
+    client = wandb.Api()
+
+    if run_id is None:
+        # process all runs
+        runs = client.runs(entity + "/" + project)
+    else:
+        try:
+            # get the run by run_id
+            run = client.run(f"{entity}/{project}/{run_id}")
+        except Exception:
+            click.echo(f"Could not find run '{entity}/{project}/{run_id}'", err=True)
+            return
+        runs = (run,)
+
+    for run in tqdm(runs, desc="Converting wandb logs"):
+        if not run.config.items():
+            continue
+        aim_run = Run(
+            repo=repo_inst,
+            system_tracking_interval=None,
+            capture_terminal_logs=False,
+            experiment=project
+        )
+        aim_run['wandb_run_id'] = run.id
+        aim_run['wandb_run_name'] = run.name
+        aim_run.description = run.notes
+
+        with TemporaryDirectory() as tmpdirname:
+            # Collect console output logs
+            console_log_filename = 'output.log'
+            console_log_file = run.file(console_log_filename)
+            try:
+                # Even though the file does not exist, a file object will be returned in zero-sized.
+                if console_log_file.size:
+                    console_log_file.download(root=tmpdirname)
+                    with open(Path(tmpdirname) / console_log_filename) as f:
+                        [aim_run.track(LogLine(line), name='logs', step=i) for i, line in enumerate(f.readlines())]
+            except Exception:
+                click.echo("Failed to track console output log.", err=True)
+
+            # TODO: Collect media files, possibly?
+
+        # Collect params & tags
+        aim_run['params'] = run.config
+        for tag in run.tags:
+            aim_run.add_tag(tag)
+
+        keys = [key for key in run.history(stream='default').keys()
+                if not key.startswith('_')]
+
+        # Collect metrics
+        for record in run.scan_history():
+            step = record.get('_step')
+            epoch = record.get('epoch')
+            timestamp = record.get('_timestamp')
+            for key in keys:
+                value = record.get(key)
+                if value is None:
+                    continue
+                try:
+                    tag, name = key.rsplit("/", 1)
+                    if "train" in tag:
+                        context = {'tag': tag, 'subset': 'train'}
+                    elif "val" in tag:
+                        context = {'tag': tag, 'subset': 'val'}
+                    elif "test" in tag:
+                        context = {'tag': tag, 'subset': 'test'}
+                    else:
+                        context = {'tag': tag}
+                except ValueError:
+                    name, context = key, {}
+                try:
+                    if timestamp:
+                        aim_run._tracker._track(value, track_time=timestamp, name=name,
+                                                step=step, epoch=epoch, context=context)
+                    else:
+                        aim_run.track(value, name=name, step=step, epoch=epoch, context=context)
+                except ValueError:
+                    click.echo(f"Type '{type(value).__name__}': artifacts are not supported yet.", err=True)
+
+        # Collect system logs
+        # NOTE: In 'system' logs, collecting sampled history cannot be avoided. (default 'samples' == 500)
+        # TODO: async history fetching for better performance
+        for record in run.history(stream='system', pandas=False, samples=1e3):
+            timestamp = record.get('_timestamp')
+            for key in record:
+                if key.startswith('_'):  # Including '_runtime', '_timestamp', '_wandb'
+                    continue
+
+                value = record.get(key)
+                if value is None:
+                    continue
+
+                name, context = _normalize_system_metric_key(key)
+                if name is None:
+                    continue
+
+                try:
+                    if timestamp:
+                        aim_run._tracker._track(value, track_time=timestamp,
+                                                name=f'{AIM_RESOURCE_METRIC_PREFIX}{name}', context=context)
+                    else:
+                        aim_run.track(value, name=f'{AIM_RESOURCE_METRIC_PREFIX}{name}', context=context)
+                except ValueError:
+                    click.echo(f"Type '{type(value).__name__}': artifacts are not supported yet.", err=True)
+
+
+def _normalize_system_metric_key(key):
+    # Remap names for being categorized as `System` in aim ui: `aim/web/ui/src/config/systemMetrics/systemMetrics.ts``
+    # {value is None} means not supported yet
+    SYSTEM_METRICS_NAME_MAP = {
+        '': {
+            'cpu': 'cpu',
+            'disk': 'disk_percent',
+            'memory': 'memory_percent',
+            'network.recv': None,
+            'network.sent': None,
+        },
+        'gpu': {
+            'gpu': 'gpu',
+            'memory': 'gpu_memory_percent',
+            'memoryAllocated': None,
+            'powerPercent': None,
+            'powerWatts': 'gpu_power_watts',
+            'temp': 'gpu_temp',
+        },
+        'proc': {
+            'cpu.threads': None,
+            'memory.availableMB': None,
+            'memory.percent': 'p_memory_percent',
+            'momory.rssMB': None,
+        }
+    }
+
+    name = re.sub(r'^system\.', '', key)
+    gpu_idx_pattern = re.compile(r'^[0-9]+\.')
+
+    # Triage & Remap name for aim ui
+    if name.startswith('gpu'):
+        name = re.sub(r'^gpu\.', '', name)
+
+        # Cut & paste gpu idx from name to context
+        gpu_idx_match = gpu_idx_pattern.search(name)
+        if gpu_idx_match:
+            gpu_idx_str = gpu_idx_match.group()
+            name = name[len(gpu_idx_str):]
+            gpu_idx = int(gpu_idx_str.rstrip('.'))
+            context = {'gpu': gpu_idx, 'tag': 'system', 'subset': 'gpu'}
+        else:
+            context = {'gpu': 'no_idx', 'tag': 'system', 'subset': 'gpu'}
+
+        normalized_name = SYSTEM_METRICS_NAME_MAP['gpu'].get(name)
+
+    elif name.startswith('proc'):
+        name = re.sub(r'^proc\.', '', name)
+        normalized_name = SYSTEM_METRICS_NAME_MAP['proc'].get(name)
+        context = {'tag': 'system', 'subset': 'proc'}
+
+    else:
+        normalized_name = SYSTEM_METRICS_NAME_MAP[''].get(name)
+        context = {'tag': 'system'}
+
+    return normalized_name, context

aim/cli/init/commands.py

@@ -0,0 +1,38 @@
+import click
+import os
+
+from aim.sdk.repo import Repo
+from aim.sdk.utils import clean_repo_path
+
+
+@click.command()
+@click.option('--repo', required=False, type=click.Path(exists=True,
+                                                        file_okay=False,
+                                                        dir_okay=True,
+                                                        writable=True))
+@click.option('-y', '--yes', is_flag=True, help='Automatically confirm prompt')
+def init(repo, yes):
+    """
+    Initializes new repository in the --repo directory.
+    Initializes new repository in the current working directory if --repo argument is not provided:
+     - Creates .aim directory & runs upgrades for structured DB
+    """
+    repo_path = clean_repo_path(repo) or os.getcwd()
+    re_init = False
+    if Repo.exists(repo_path):
+        if yes:
+            re_init = True
+        else:
+            re_init = click.confirm('Aim repository is already initialized. '
+                                    'Do you want to re-initialize to empty Aim repository?')
+        if not re_init:
+            return
+        # Clear old repo
+        Repo.rm(repo_path)
+
+    repo = Repo.from_path(repo_path, init=True)
+    if re_init:
+        click.echo(
+            'Re-initialized empty Aim repository at {}'.format(repo.root_path))
+    else:
+        click.echo('Initialized a new Aim repository at {}'.format(repo.root_path))

aim/cli/manager/manager.py

@@ -0,0 +1,128 @@
+import subprocess
+import time
+import enum
+
+from aim.cli.configs import VERSION_NAME, UP_NAME
+from aim.__version__ import __version__
+
+# Error message prefix for aim commands
+ERROR_MSG_PREFIX = 'Error:'
+
+
+class ManagerActionStatuses(enum.Enum):
+    Failed = 1,
+    Succeed = 2,
+
+
+class ManagerActionResult:
+    """
+        Object returned by manager action
+        If status is ManagerActionStatuses.Failed the info dict should have a message property
+        If status is ManagerActionStatuses.Succeed the info dict should have required properties for the specific action
+        @TODO add type checking for info fields
+    """
+    def __init__(self, status: ManagerActionStatuses, info: dict = None):
+        self.status = status
+        self.info = info
+
+
+def run_up(args):
+    def check_startup_success():
+        import requests
+        server_path = 'http://{}:{}{}'.format(args['--host'], args['--port'], args['--base-path'])
+        status_api = f'{server_path}/api/projects/status'
+        retry_count = 5
+        sleep_interval = 1
+        for _ in range(retry_count):
+            try:
+                response = requests.get(status_api)
+                if response.status_code == 200:
+                    return True
+            except Exception:
+                pass
+            sleep_interval += 1
+            time.sleep(sleep_interval)
+
+        return False
+
+    args_list = ['--log-level=error']
+    for p in args.keys():
+        if p != '--proxy-url':
+            args_list.append(p + '=' + args[p])
+
+    child_process = subprocess.Popen(
+        ['aim', UP_NAME] + args_list + ['--force-init'],
+        stderr=subprocess.PIPE,
+        stdout=subprocess.PIPE
+    )
+    # Runs `aim up <args>` command
+    info = {
+        'port': args['--port'],
+        'host': 'http://' + args['--host']
+    }
+
+    if check_startup_success():
+        return ManagerActionResult(
+            ManagerActionStatuses.Succeed,
+            info
+        )
+
+    for line in child_process.stderr:
+        if ERROR_MSG_PREFIX in line.decode():
+            return ManagerActionResult(
+                ManagerActionStatuses.Failed,
+                {'message': line.decode()}
+            )
+
+    return ManagerActionResult(
+        ManagerActionStatuses.Failed,
+        {
+            'message': '\nPerhaps this is a bug from aim side.'
+                       '\nPlease open an issue https://github.com/aimhubio/aim/issues.'
+        }
+    )
+
+
+def run_version(args):
+    """Aim cli command runner to get aim current version
+    Returns:
+         the current version of aim
+    """
+    return ManagerActionResult(
+        ManagerActionStatuses.Succeed,
+        {'version': __version__}
+    )
+
+
+# command runners dict
+COMMANDS = {
+    UP_NAME: run_up,
+    VERSION_NAME: run_version,
+    # Perhaps it will need to implement the help command as well
+}
+
+
+def __get_command_runner(command):
+    """Returns the runner function of the specified command"""
+    return COMMANDS[command]
+
+
+def validate_command(command):
+    """Validate command existence
+    """
+    return command in (VERSION_NAME, UP_NAME)
+
+
+def run_process(command, args):
+    """
+    Returns None if command is not exists
+            call the corresponding runner function to  execute aim command `aim <command> <args>`
+    """
+    if validate_command(command) is not True:
+        return ManagerActionResult(
+            ManagerActionStatuses.Failed,
+            {'message': 'Invalid operation'}
+        )
+
+    run = __get_command_runner(command)
+    return run(args)

aim/cli/reindex/commands.py

@@ -0,0 +1,18 @@
+import click
+
+
+@click.command()
+@click.option('--repo', required=False, type=click.Path(exists=True,
+                                                        file_okay=False,
+                                                        dir_okay=True,
+                                                        writable=True))
+@click.option('--finalize-only', required=False, is_flag=True, default=False)
+def reindex(repo, finalize_only):
+    """
+    Process runs left in 'in progress' state.
+    """
+    from aim.utils.deprecation import deprecation_warning
+
+    deprecation_warning(remove_version='3.16', msg='`aim reindex` is deprecated! '
+                                                   'Use `aim runs close` command instead.')
+    return

aim/cli/runs/commands.py

@@ -0,0 +1,163 @@
+import click
+import os
+import tqdm
+
+from multiprocessing.pool import ThreadPool
+from psutil import cpu_count
+
+from aim.cli.runs.utils import match_runs, make_zip_archive, upload_repo_runs
+from aim.sdk.repo import Repo
+
+
+@click.group()
+@click.option('--repo', required=False,
+              default=os.getcwd(),
+              type=str)
+@click.pass_context
+def runs(ctx, repo):
+    """Manage runs in aim repository."""
+    ctx.ensure_object(dict)
+    ctx.obj['repo'] = repo
+
+
+@runs.command(name='ls')
+@click.pass_context
+def list_runs(ctx):
+    """List Runs available in Repo."""
+    repo_path = ctx.obj['repo']
+    if not Repo.is_remote_path(repo_path):
+        if not Repo.exists(repo_path):
+            click.echo(f'\'{repo_path}\' is not a valid aim repo.')
+            exit(1)
+
+    repo = Repo.from_path(repo_path)
+    run_hashes = repo.list_all_runs()
+
+    click.echo('\t'.join(run_hashes))
+    click.echo(f'Total {len(run_hashes)} runs.')
+
+
+@runs.command(name='rm')
+@click.argument('hashes', nargs=-1, type=str)
+@click.pass_context
+@click.option('-y', '--yes', is_flag=True, help='Automatically confirm prompt')
+def remove_runs(ctx, hashes, yes):
+    """Remove Run data for given run hashes."""
+    if len(hashes) == 0:
+        click.echo('Please specify at least one Run to delete.')
+        exit(1)
+    repo_path = ctx.obj['repo']
+    repo = Repo.from_path(repo_path)
+
+    matched_hashes = match_runs(repo, hashes)
+    if yes:
+        confirmed = True
+    else:
+        confirmed = click.confirm(f'This command will permanently delete {len(matched_hashes)} runs from aim repo '
+                                  f'located at \'{repo_path}\'. Do you want to proceed?')
+    if not confirmed:
+        return
+
+    success, remaining_runs = repo.delete_runs(matched_hashes)
+    if success:
+        click.echo(f'Successfully deleted {len(matched_hashes)} runs.')
+    else:
+        click.echo('Something went wrong while deleting runs. Remaining runs are:', err=True)
+        click.secho('\t'.join(remaining_runs), fg='yellow')
+
+
+@runs.command(name='cp')
+@click.option('--destination', required=True, type=str)
+@click.argument('hashes', nargs=-1, type=str)
+@click.pass_context
+def copy_runs(ctx, destination, hashes):
+    """Copy Run data for given run hashes to destination Repo."""
+    if len(hashes) == 0:
+        click.echo('Please specify at least one Run to copy.')
+        exit(1)
+    source = ctx.obj['repo']
+    source_repo = Repo.from_path(source)
+    destination_repo = Repo.from_path(destination)
+
+    matched_hashes = match_runs(source_repo, hashes)
+    success, remaining_runs = source_repo.copy_runs(matched_hashes, destination_repo)
+    if success:
+        click.echo(f'Successfully copied {len(matched_hashes)} runs.')
+    else:
+        click.echo('Something went wrong while copying runs. Remaining runs are:', err=True)
+        click.secho('\t'.join(remaining_runs), fg='yellow')
+
+
+@runs.command(name='mv')
+@click.option('--destination', required=True,
+              type=str)
+@click.argument('hashes', nargs=-1, type=str)
+@click.pass_context
+def move_runs(ctx, destination, hashes):
+    """Move Run data for given run hashes to destination Repo."""
+    if len(hashes) == 0:
+        click.echo('Please specify at least one Run to move.')
+        exit(1)
+    source = ctx.obj['repo']
+    source_repo = Repo.from_path(source)
+    destination_repo = Repo.from_path(destination)
+
+    matched_hashes = match_runs(source_repo, hashes)
+
+    success, remaining_runs = source_repo.move_runs(matched_hashes, destination_repo)
+    if success:
+        click.echo(f'Successfully moved {len(matched_hashes)} runs.')
+    else:
+        click.echo('Something went wrong while moving runs. Remaining runs are:', err=True)
+        click.secho('\t'.join(remaining_runs), fg='yellow')
+
+
+@runs.command(name='upload')
+@click.argument('bucket', nargs=1, type=str)
+@click.pass_context
+def upload_runs(ctx, bucket):
+    """Upload Repo backup to the given S3 bucket."""
+    repo_path = ctx.obj['repo']
+    if not Repo.exists(repo_path):
+        click.echo(f'\'{repo_path}\' is not a valid aim repo.')
+        exit(1)
+
+    zip_buffer = make_zip_archive(repo_path)
+    zip_buffer.seek(0)
+
+    success, uploaded_zip_file_name = upload_repo_runs(zip_buffer, bucket)
+    if success:
+        click.echo(f'Successfully uploaded runs in {uploaded_zip_file_name}.')
+    else:
+        click.echo(f'The storage backup failed because of the following error: {uploaded_zip_file_name}.')
+
+
+@runs.command(name='close')
+@click.argument('hashes', nargs=-1, type=str)
+@click.pass_context
+@click.option('-y', '--yes', is_flag=True, help='Automatically confirm prompt')
+def close_runs(ctx, hashes, yes):
+    """Close failed/stalled Runs."""
+    repo_path = ctx.obj['repo']
+    repo = Repo.from_path(repo_path)
+
+    if len(hashes) == 0:
+        click.echo('Please specify at least one Run to close.')
+        exit(1)
+
+    click.secho(f'This command will forcefully close {len(hashes)} Runs from Aim Repo \'{repo_path}\'. '
+                f'Please make sure Runs are not active. Data corruption may occur otherwise.')
+    if yes:
+        confirmed = True
+    else:
+        confirmed = click.confirm('Do you want to proceed?')
+    if not confirmed:
+        return
+
+    pool = ThreadPool(cpu_count(logical=False))
+
+    for _ in tqdm.tqdm(
+            pool.imap_unordered(repo._close_run, hashes),
+            desc='Closing runs',
+            total=len(hashes)):
+        pass

aim/cli/runs/utils.py

@@ -0,0 +1,68 @@
+import fnmatch
+import os
+
+from typing import List, Tuple, TYPE_CHECKING
+import io
+import zipfile
+from datetime import datetime
+
+if TYPE_CHECKING:
+    from aim.sdk.repo import Repo
+
+
+def match_runs(repo: 'Repo', hashes: List[str], lookup_dir: str = None) -> List[str]:
+    matched_hashes = set()
+    all_run_hashes = None
+    for run_hash in hashes:
+        if '*' in run_hash:
+            expr = run_hash  # for the sake of readability
+            # avoiding multiple or unnecessary list_runs() calls
+            if not all_run_hashes:
+                all_run_hashes = repo.list_all_runs()
+            if expr == '*':
+                return all_run_hashes
+            # update the matches set with current expression matches
+            matched_hashes.update(fnmatch.filter(all_run_hashes, expr))
+        else:
+            matched_hashes.add(run_hash)
+
+    return list(matched_hashes)
+
+
+def make_zip_archive(repo_path: str) -> io.BytesIO:
+    aim_dir = os.path.join(repo_path, '.aim')
+    zip_buf = io.BytesIO()
+    zipf = zipfile.ZipFile(zip_buf, 'w', zipfile.ZIP_DEFLATED)
+    len_dir_path = len(aim_dir)
+    for root, _, files in os.walk(aim_dir):
+        for file in files:
+            file_path = os.path.join(root, file)
+            zipf.write(file_path, file_path[len_dir_path:])
+    zipf.close()
+    return zip_buf
+
+
+def upload_repo_runs(buffer: io.BytesIO, bucket_name: str) -> Tuple[bool, str]:
+    try:
+        import boto3
+    except ImportError:
+        raise RuntimeError(
+            'This command requires \'boto3\' to be installed. '
+            'Please install it with command: \n pip install boto3'
+        )
+
+    try:
+        s3_client = boto3.client('s3')
+        buckets = s3_client.list_buckets()
+        bucket_names = []
+        for bucket in buckets['Buckets']:
+            bucket_names.append(bucket['Name'])
+
+        if bucket_name not in bucket_names:
+            s3_client.create_bucket(Bucket=bucket_name)
+
+        key = f'aim-{datetime.now().strftime("%Y-%m-%d-%H-%M-%S")}.zip'
+        s3_client.upload_fileobj(buffer, bucket_name, key)
+        return True, key
+    except Exception as e:
+        return False, e

aim/cli/server/commands.py

@@ -0,0 +1,96 @@
+import os
+import click
+
+from aim.sdk.repo import Repo
+from aim.sdk.utils import clean_repo_path
+from aim.cli.utils import (
+    ShellCommandException,
+    build_uvicorn_command,
+    exec_cmd,
+    get_free_port_num,
+    get_repo_instance,
+    set_log_level
+)
+from aim.ext.transport.config import (
+    AIM_SERVER_DEFAULT_PORT,
+    AIM_SERVER_DEFAULT_HOST,
+    AIM_SERVER_MOUNTED_REPO_PATH,
+    AIM_SERVER_BASE_PATH
+)
+from aim.web.configs import AIM_ENV_MODE_KEY
+
+
+@click.command('server')
+@click.option('-h', '--host', default=AIM_SERVER_DEFAULT_HOST, type=str)
+@click.option('-p', '--port', default=AIM_SERVER_DEFAULT_PORT, type=int)
+@click.option('--repo', required=False, default=os.getcwd(), type=click.Path(exists=True,
+                                                                             file_okay=False,
+                                                                             dir_okay=True,
+                                                                             writable=True))
+@click.option('--ssl-keyfile', required=False, type=click.Path(exists=True,
+                                                               file_okay=True,
+                                                               dir_okay=False,
+                                                               readable=True))
+@click.option('--ssl-certfile', required=False, type=click.Path(exists=True,
+                                                                file_okay=True,
+                                                                dir_okay=False,
+                                                                readable=True))
+@click.option('--base-path', required=False, default='', type=str)
+@click.option('--log-level', required=False, default='', type=str)
+@click.option('--dev', is_flag=True, default=False)
+@click.option('-y', '--yes', is_flag=True, help='Automatically confirm prompt')
+def server(host, port,
+           repo, ssl_keyfile, ssl_certfile,
+           base_path, log_level, dev, yes):
+    """
+    Starts the Aim remote tracking server for real-time logging.
+
+    The Aim tracking server facilitates real-time logging of experiments
+    from remote locations. This command launches the server with specified
+    configurations, including host, port, and associated repository.
+
+    Like the UI, the server can also run in production or development mode.
+    """
+    if dev:
+        os.environ[AIM_ENV_MODE_KEY] = 'dev'
+        log_level = log_level or 'debug'
+    else:
+        os.environ[AIM_ENV_MODE_KEY] = 'prod'
+
+    if log_level:
+        set_log_level(log_level)
+
+    if base_path:
+        if base_path.endswith('/'):
+            base_path = base_path[:-1]
+        if not base_path.startswith('/'):
+            base_path = f'/{base_path}'
+        os.environ[AIM_SERVER_BASE_PATH] = base_path
+
+    if port == 0:
+        try:
+            port = get_free_port_num()
+        except Exception:
+            pass
+
+    repo_path = clean_repo_path(repo) or Repo.default_repo_path()
+    repo_inst = get_repo_instance(repo_path, yes)
+
+    if not repo_inst:
+        return
+
+    os.environ[AIM_SERVER_MOUNTED_REPO_PATH] = repo_inst.path
+
+    click.secho('Running Aim Server on repo `{}`'.format(repo), fg='yellow')
+    click.echo('Server is mounted on aim://{}:{}'.format(host, port), err=True)
+    click.echo('Press Ctrl+C to exit')
+
+    try:
+        cmd = build_uvicorn_command('aim.ext.transport.run:app',
+                                    host=host, port=port,
+                                    ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile, log_level=log_level)
+        exec_cmd(cmd, stream_output=True)
+    except ShellCommandException:
+        click.echo('Failed to run Aim Tracking Server. '
+                   'Please see the logs above for details.')
+        exit(1)

aim/cli/storage/commands.py

@@ -0,0 +1,133 @@
+import click
+import os
+from tqdm import tqdm
+
+from aim.cli.runs.utils import match_runs
+
+from aim.sdk.maintenance_run import MaintenanceRun as Run
+from aim.sdk.utils import backup_run, restore_run_backup
+from aim.sdk.repo import Repo
+from aim.sdk.index_manager import RepoIndexManager
+
+
+@click.group()
+@click.option('--repo', required=False,
+              default=os.getcwd(),
+              type=str)
+@click.pass_context
+def storage(ctx, repo):
+    """Manage aim repository data & format updates."""
+    ctx.ensure_object(dict)
+    ctx.obj['repo'] = repo
+
+
+@storage.group()
+@click.pass_context
+def upgrade(ctx):
+    """Update Runs data for given run hashes to use new format."""
+    pass
+
+
+@upgrade.command(name='3.11+')
+@click.argument('hashes', nargs=-1, type=str)
+@click.pass_context
+@click.option('-y', '--yes', is_flag=True, help='Automatically confirm prompt')
+def to_3_11(ctx, hashes, yes):
+    """Optimize Runs Metrics data for read access."""
+    if len(hashes) == 0:
+        click.echo('Please specify at least one Run to update.')
+        exit(1)
+    repo_path = ctx.obj['repo']
+    repo = Repo.from_path(repo_path)
+
+    matched_hashes = match_runs(repo, hashes)
+    remaining_runs = []
+    if yes:
+        confirmed = True
+    else:
+        confirmed = click.confirm(f'This command will optimize the metrics data for {len(matched_hashes)} '
+                                  f'runs from aim repo located at \'{repo_path}\'. This process might take a while. '
+                                  f'Do you want to proceed?')
+    if not confirmed:
+        return
+
+    index_manager = RepoIndexManager.get_index_manager(repo)
+    for run_hash in tqdm(matched_hashes):
+        try:
+            run = Run(run_hash, repo=repo)
+            if run.check_metrics_version():
+                backup_run(run)
+                run.update_metrics()
+                index_manager.index(run_hash)
+            else:
+                click.echo(f'Run {run.hash} is already up to date. Skipping')
+        except Exception:
+            remaining_runs.append(run_hash)
+
+    if not remaining_runs:
+        click.echo('Finished optimizing metric data.')
+    else:
+        click.echo('Finished optimizing metric data. The following runs were skipped:')
+        click.secho(' '.join(remaining_runs), fg='yellow')
+    click.echo('In case of any issues the following command can be used to restore data:')
+    click.secho(f'aim storage --repo {repo.root_path} restore \'*\'', fg='yellow')
+
+
+@storage.command(name='restore')
+@click.argument('hashes', nargs=-1, type=str)
+@click.pass_context
+@click.option('-y', '--yes', is_flag=True, help='Automatically confirm prompt')
+def restore_runs(ctx, hashes, yes):
+    """Rollback Runs data for given run hashes to the previous metric format. """
+    if len(hashes) == 0:
+        click.echo('Please specify at least one Run to delete.')
+        exit(1)
+    repo_path = ctx.obj['repo']
+    repo = Repo.from_path(repo_path)
+
+    matched_hashes = match_runs(repo, hashes, lookup_dir='bcp')
+    if yes:
+        confirmed = True
+    else:
+        confirmed = click.confirm(f'This command will restore {len(matched_hashes)} runs from aim repo '
+                                  f'located at \'{repo_path}\'. Do you want to proceed?')
+    if not confirmed:
+        return
+
+    remaining_runs = []
+    index_manager = RepoIndexManager.get_index_manager(repo)
+    for run_hash in tqdm(matched_hashes):
+        try:
+            restore_run_backup(repo, run_hash)
+            index_manager.index(run_hash)
+        except Exception as e:
+            click.echo(f'Error while trying to restore run \'{run_hash}\'. {str(e)}.', err=True)
+            remaining_runs.append(run_hash)
+
+    if not remaining_runs:
+        click.echo(f'Successfully restored {len(matched_hashes)} runs.')
+    else:
+        click.echo('Something went wrong while restoring runs. Remaining runs are:', err=True)
+        click.secho('\t'.join(remaining_runs), fg='yellow')
+
+
+@storage.command(name='prune')
+@click.pass_context
+def prune(ctx):
+    """Remove dangling/orphan params/sequences with no referring runs."""
+
+    repo_path = ctx.obj['repo']
+    repo = Repo.from_path(repo_path)
+    repo.prune()
+
+
+@storage.command('reindex')
+@click.option('--finalize-only', required=False, is_flag=True, default=False)
+@click.pass_context
+def reindex(ctx, finalize_only):
+    """ Process runs left in 'in progress' state. """
+    from aim.utils.deprecation import deprecation_warning
+
+    deprecation_warning(remove_version='3.16', msg='`aim storage reindex` is deprecated! '
+                                                   'Use `aim runs close` command instead.')
+    return

aim/cli/up/commands.py

@@ -0,0 +1,127 @@
+import os
+import click
+
+from aim.cli.utils import (
+    ShellCommandException,
+    build_db_upgrade_command,
+    build_uvicorn_command,
+    exec_cmd,
+    get_free_port_num,
+    get_repo_instance,
+    set_log_level
+)
+from aim.web.configs import (
+    AIM_ENV_MODE_KEY,
+    AIM_TF_LOGS_PATH_KEY,
+    AIM_UI_BASE_PATH,
+    AIM_UI_DEFAULT_HOST,
+    AIM_UI_DEFAULT_PORT,
+    AIM_UI_MOUNTED_REPO_PATH,
+    AIM_PROXY_URL,
+    AIM_PROFILER_KEY
+)
+from aim.sdk.repo import Repo
+from aim.sdk.utils import clean_repo_path
+
+
+@click.command()
+@click.option('-h', '--host', default=AIM_UI_DEFAULT_HOST, type=str)
+@click.option('-p', '--port', default=AIM_UI_DEFAULT_PORT, type=int)
+@click.option('-w', '--workers', default=1, type=int)
+@click.option('--uds', required=False, type=click.Path(exists=False,
+                                                       file_okay=True,
+                                                       dir_okay=False,
+                                                       readable=True))
+@click.option('--repo', required=False, type=click.Path(exists=True,
+                                                        file_okay=False,
+                                                        dir_okay=True,
+                                                        writable=True))
+@click.option('--tf_logs', type=click.Path(exists=True, readable=True))
+@click.option('--dev', is_flag=True, default=False)
+@click.option('--ssl-keyfile', required=False, type=click.Path(exists=True,
+                                                               file_okay=True,
+                                                               dir_okay=False,
+                                                               readable=True))
+@click.option('--ssl-certfile', required=False, type=click.Path(exists=True,
+                                                                file_okay=True,
+                                                                dir_okay=False,
+                                                                readable=True))
+@click.option('--base-path', required=False, default='', type=str)
+@click.option('--profiler', is_flag=True, default=False)
+@click.option('--log-level', required=False, default='', type=str)
+@click.option('-y', '--yes', is_flag=True, help='Automatically confirm prompt')
+def up(dev, host, port, workers, uds,
+       repo, tf_logs,
+       ssl_keyfile, ssl_certfile,
+       base_path, profiler,
+       log_level, yes):
+    if dev:
+        os.environ[AIM_ENV_MODE_KEY] = 'dev'
+        log_level = log_level or 'debug'
+    else:
+        os.environ[AIM_ENV_MODE_KEY] = 'prod'
+
+    if log_level:
+        set_log_level(log_level)
+
+    if base_path:
+        # process `base_path` as ui requires leading slash
+        if base_path.endswith('/'):
+            base_path = base_path[:-1]
+        if base_path and not base_path.startswith('/'):
+            base_path = f'/{base_path}'
+        os.environ[AIM_UI_BASE_PATH] = base_path
+
+    repo_path = clean_repo_path(repo) or Repo.default_repo_path()
+    repo_inst = get_repo_instance(repo_path, yes)
+
+    if not repo_inst:
+        return
+
+    os.environ[AIM_UI_MOUNTED_REPO_PATH] = repo_inst.path
+
+    if tf_logs:
+        os.environ[AIM_TF_LOGS_PATH_KEY] = tf_logs
+
+    try:
+        db_cmd = build_db_upgrade_command()
+        exec_cmd(db_cmd, stream_output=True)
+    except ShellCommandException:
+        click.echo('Failed to initialize Aim DB. '
+                   'Please see the logs above for details.')
+        return
+
+    if port == 0:
+        try:
+            port = get_free_port_num()
+        except Exception:
+            pass
+
+    click.echo(click.style('Running Aim UI on repo `{}`'.format(repo_inst), fg='yellow'))
+
+    if uds:
+        click.echo('Aim UI running on {}'.format(uds))
+    else:
+        scheme = 'https' if ssl_keyfile or ssl_certfile else 'http'
+        click.echo('Open {}://{}:{}{}'.format(scheme, host, port, base_path), err=True)
+
+    proxy_url = os.environ.get(AIM_PROXY_URL)
+    if proxy_url:
+        click.echo(f'Proxy {proxy_url}{base_path}/')
+
+    click.echo('Press Ctrl+C to exit')
+
+    if profiler:
+        os.environ[AIM_PROFILER_KEY] = '1'
+
+    try:
+        server_cmd = build_uvicorn_command(
+            'aim.web.run:app',
+            host=host, port=port,
+            workers=workers, uds=uds,
+            ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile,
+            log_level=log_level)
+        exec_cmd(server_cmd, stream_output=True)
+    except ShellCommandException:
+        click.echo('Failed to run Aim UI. Please see the logs above for details.')
+        return

aim/cli/utils.py

@@ -0,0 +1,153 @@
+import click
+import os
+import sys
+import logging
+import subprocess
+
+from typing import Optional
+
+from aim.sdk.repo import Repo, RepoStatus
+from aim.web.configs import AIM_ENV_MODE_KEY
+from aim.web.configs import AIM_LOG_LEVEL_KEY
+
+
+class ShellCommandException(Exception):
+    pass
+
+
+def exec_cmd(
+    cmd, throw_on_error=True, env=None, stream_output=False, cwd=None, cmd_stdin=None, **kwargs
+):
+    """
+    Runs a command as a child process.
+    A convenience wrapper for running a command from a Python script.
+    Keyword arguments:
+    cmd -- the command to run, as a list of strings
+    throw_on_error -- if true, raises an Exception if the exit code of the program is nonzero
+    env -- additional environment variables to be defined when running the child process
+    cwd -- working directory for child process
+    stream_output -- if true, does not capture standard output and error; if false, captures these
+      streams and returns them
+    cmd_stdin -- if specified, passes the specified string as stdin to the child process.
+    Note on the return value: If stream_output is true, then only the exit code is returned. If
+    stream_output is false, then a tuple of the exit code, standard output and standard error is
+    returned.
+    """
+    cmd_env = os.environ.copy()
+    if env:
+        cmd_env.update(env)
+    if stream_output:
+        child = subprocess.Popen(
+            cmd, env=cmd_env, cwd=cwd, universal_newlines=True, stdin=subprocess.PIPE, **kwargs
+        )
+        child.communicate(cmd_stdin)
+        exit_code = child.wait()
+        if throw_on_error and exit_code != 0:
+            raise ShellCommandException("Non-zero exitcode: %s" % (exit_code))
+        return exit_code
+    else:
+        child = subprocess.Popen(
+            cmd,
+            env=cmd_env,
+            stdout=subprocess.PIPE,
+            stdin=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            cwd=cwd,
+            universal_newlines=True,
+            **kwargs
+        )
+        (stdout, stderr) = child.communicate(cmd_stdin)
+        exit_code = child.wait()
+        if throw_on_error and exit_code != 0:
+            raise ShellCommandException(
+                "Non-zero exit code: %s\n\nSTDOUT:\n%s\n\nSTDERR:%s" % (exit_code, stdout, stderr)
+            )
+        return exit_code, stdout, stderr
+
+
+def set_log_level(log_level):
+    numeric_level = getattr(logging, log_level.upper(), None)
+    if not isinstance(numeric_level, int):
+        raise ValueError('Invalid log level: %s' % log_level)
+    os.environ[AIM_LOG_LEVEL_KEY] = str(numeric_level)
+    logging.basicConfig(level=numeric_level)
+
+
+def build_db_upgrade_command():
+    from aim import web
+    web_dir = os.path.dirname(web.__file__)
+    migrations_dir = os.path.join(web_dir, 'migrations')
+    if os.getenv(AIM_ENV_MODE_KEY, 'prod') == 'prod':
+        ini_file = os.path.join(migrations_dir, 'alembic.ini')
+    else:
+        ini_file = os.path.join(migrations_dir, 'alembic_dev.ini')
+    return [sys.executable, '-m', 'alembic', '-c', ini_file, 'upgrade', 'head']
+
+
+def build_uvicorn_command(app,
+                          host='0.0.0.0',
+                          port=0,
+                          workers=1,
+                          uds=None,
+                          ssl_keyfile=None,
+                          ssl_certfile=None,
+                          log_level='warning',
+                          ):
+    cmd = [sys.executable, '-m', 'uvicorn',
+           '--host', host, '--port', f'{port}',
+           '--workers', f'{workers}']
+
+    if os.getenv(AIM_ENV_MODE_KEY, 'prod') == 'prod':
+        log_level = log_level or 'error'
+    else:
+        import aim
+        cmd += ['--reload', '--reload-dir', os.path.dirname(aim.__file__)]
+        log_level = log_level or 'debug'
+
+    if uds:
+        cmd += ['--uds', uds]
+    if ssl_keyfile:
+        cmd += ['--ssl-keyfile', ssl_keyfile]
+    if ssl_certfile:
+        cmd += ['--ssl-certfile', ssl_certfile]
+    cmd += ['--log-level', log_level.lower()]
+    cmd += [app]
+    return cmd
+
+
+def get_free_port_num():
+    import socket
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.bind(('', 0))
+    port_num = s.getsockname()[1]
+    s.close()
+    return port_num
+
+
+def get_repo_instance(repo_path: str, yes: bool) -> Optional['Repo']:
+    repo_status = Repo.check_repo_status(repo_path)
+    if repo_status == RepoStatus.MISSING:
+        if yes:
+            init_repo = True
+        else:
+            init_repo = click.confirm(f'\'{repo_path}\' is not a valid Aim repository. Do you want to initialize it?')
+        if not init_repo:
+            click.echo('To initialize repo please run the following command:')
+            click.secho('aim init', fg='yellow')
+            return
+        repo_inst = Repo.from_path(repo_path, init=True)
+    elif repo_status == RepoStatus.UPDATE_REQUIRED:
+        if yes:
+            reinit_repo = True
+        else:
+            reinit_repo = click.confirm('Found non-empty \'.aim\' directory. Would you like to overwrite it?')
+        if not reinit_repo:
+            click.echo('To re-initialize repo please run the following command:')
+            click.secho('aim init', fg='yellow')
+            return
+        Repo.rm(repo_path)
+        repo_inst = Repo.from_path(repo_path, init=True)
+    else:
+        repo_inst = Repo.from_path(repo_path)
+
+    return repo_inst

aim/cli/version/commands.py

@@ -0,0 +1,8 @@
+import click
+
+from aim.__version__ import __version__ as aim_version
+
+
+@click.command()
+def version():
+    click.echo('Aim v{}'.format(aim_version))

aim/cli/watcher_cli.py

@@ -0,0 +1,296 @@
+import logging
+
+import click
+import uuid
+
+from click import core
+from collections import OrderedDict
+from typing import Optional, Mapping
+
+from aim.ext.notifier import get_config
+from aim.ext.notifier.utils import has_watcher_config, set_default_config
+from aim.sdk.run_status_watcher import RunStatusWatcher
+from aim.sdk.repo import Repo
+
+
+core._verify_python3_env = lambda: None
+DEFAULT_MESSAGE_TEMPLATE = "❗️ Something wrong with Run '{run.hash}'. Please check. ❗️"
+MESSAGE_PROMPT = "Stuck Runs notification message"
+
+
+class OrderedGroup(click.Group):
+    def __init__(self, name: Optional[str] = None, commands: Optional[Mapping[str, click.Command]] = None, **kwargs):
+        super(OrderedGroup, self).__init__(name, commands, **kwargs)
+        #: the registered subcommands by their exported names.
+        self.commands = commands or OrderedDict()
+
+    def list_commands(self, ctx: click.Context) -> Mapping[str, click.Command]:
+        return self.commands
+
+
+def check_configuration(ctx: click.Context, repo: Repo) -> bool:
+    if not has_watcher_config(repo.path):
+        click.echo(f'Repo \'{repo.path}\' has no configured notifiers.')
+        if click.confirm('Would you like to configure notifiers?', default=True):
+            ctx.invoke(add_config)
+        elif click.confirm('Would you like to use default configuration?'):
+            set_default_config(repo.path)
+        else:
+            return False
+    return True
+
+
+def dump_notifier_config(cfg: dict):
+    lines = ['', '', f'Type: {cfg["type"]}']
+    for arg_name, value in cfg['arguments'].items():
+        printable_name = arg_name.replace('_', ' ').capitalize()
+        lines.append(f'{printable_name}: {value}')
+    lines.append('--------')
+    click.echo('\n'.join(lines))
+
+
+@click.group()
+@click.option('--repo',
+              required=False,
+              help='Aim Repo to check Run statuses.',
+              type=click.Path(exists=True,
+                              file_okay=False,
+                              dir_okay=True,
+                              writable=True))
+@click.pass_context
+def cli_entry_point(ctx, repo):
+    """Service for detecting and reporting training Run failures."""
+    repo_path = repo or Repo.default_repo_path()
+    repo = Repo.from_path(repo_path)
+
+    ctx.ensure_object(dict)
+    ctx.obj['repo'] = repo
+    ctx.obj['config'] = get_config(repo.path)
+
+
+@cli_entry_point.command(name='start')
+@click.pass_context
+def start_watcher(ctx):
+    """Start watcher service to monitor and report stuck/failed Runs."""
+    repo = ctx.obj['repo']
+    if check_configuration(ctx, repo):
+        watcher = RunStatusWatcher(repo)
+        click.secho(f'Starting Aim watcher for repo \'{repo.path}\'...', fg='yellow')
+        click.echo('Press Ctrl+C to exit')
+
+        watcher.start_watcher()
+
+
+@cli_entry_point.group(cls=OrderedGroup, name='notifiers')
+@click.pass_context
+def config_notifiers(ctx):
+    """Configure how notifications should be received."""
+    pass
+
+
+@click.command(name='dump', hidden=True)
+@click.pass_context
+def dump_config(ctx):
+    """Dump notifier configuration file."""
+    cfg = ctx.obj['config']
+    if not cfg.exists():
+        repo = ctx.obj['repo']
+        click.echo(f'Cannot find notifier configuration for Repo \'{repo.path}\'.')
+        return
+
+    click.echo(cfg.dump())
+
+
+@click.command(name='list')
+@click.pass_context
+def list_config(ctx):
+    """List available notifiers."""
+    cfg = ctx.obj['config']
+    if not cfg.exists():
+        repo = ctx.obj['repo']
+        click.echo(f'Cannot find notifier configuration for Repo \'{repo.path}\'.')
+        return
+
+    click.echo("{:<40} {:<10} {:<10}".format('NOTIFIER ID', 'TYPE', 'STATUS'))
+    for notifier in cfg.notifiers.values():
+        click.echo("{:<40} {:<10} {:<10}".format(notifier['id'], notifier['type'], notifier['status']))
+
+
+@click.command(name='get-log-level')
+@click.pass_context
+def get_log_level(ctx):
+    """Get Log Notifications level."""
+    cfg = ctx.obj['config']
+    if not cfg.exists():
+        repo = ctx.obj['repo']
+        click.echo(f'Cannot find notifier configuration for Repo \'{repo.path}\'.')
+        return
+
+    click.echo(f'Log level: {logging.getLevelName(cfg.log_level)}')
+
+
+def get_level_names():
+    available_levels = (logging.CRITICAL, logging.ERROR, logging.WARNING, logging.INFO, logging.DEBUG)
+    return tuple(map(logging.getLevelName, available_levels))
+
+
+@click.command(name='set-log-level')
+@click.argument('level', required=True, type=click.Choice(get_level_names()))
+@click.pass_context
+def set_log_level(ctx, level):
+    """Set Log Notifications level to <level>."""
+    cfg = ctx.obj['config']
+    if not cfg.exists():
+        repo = ctx.obj['repo']
+        click.echo(f'Cannot find notifier configuration for Repo \'{repo.path}\'.')
+        return
+
+    cfg.log_level = getattr(logging, level)
+    cfg.save()
+
+
+@click.group(name='add', invoke_without_command=True)
+@click.pass_context
+def add_config(ctx):
+    """Add a new notifier configuration (slack, workplace, etc.)."""
+    if ctx.invoked_subcommand is None:
+        add_new = True
+        while add_new:
+            sub_commands = ctx.command.list_commands(ctx)
+            choice = click.prompt('Select notifier type to add:', show_choices=True, type=click.Choice(sub_commands))
+            sub_cmd = ctx.command.get_command(ctx, choice)
+            extra_args = {}
+            for param in sub_cmd.params:
+                if isinstance(param, click.Option) and param.prompt:
+                    extra_args[param.name] = param.prompt_for_value(ctx)
+            ctx.invoke(sub_cmd, **extra_args)
+            add_new = click.confirm('Would you like to add another notifier?')
+
+
+@click.command(name='remove')
+@click.argument('notifier-id', required=True, type=str)
+@click.pass_context
+def remove_config(ctx, notifier_id):
+    """Remove notifier configuration from the list."""
+    cfg = ctx.obj['config']
+    if not cfg.exists():
+        repo = ctx.obj['repo']
+        click.echo(f'Cannot find notifier configuration for Repo \'{repo.path}\'.')
+        return
+
+    if notifier_id in cfg.notifiers:
+        dump_notifier_config(cfg.get(notifier_id))
+        click.confirm('Remove notifier configuration above?', abort=True)
+        cfg.remove(notifier_id)
+        cfg.save()
+    else:
+        click.echo(f'No notifier with id {notifier_id} found.')
+
+
+@click.command(name='enable')
+@click.argument('notifier-id', required=True, type=str)
+@click.pass_context
+def enable_config(ctx, notifier_id):
+    """Start receiving notifications from given notifier."""
+    cfg = ctx.obj['config']
+    if not cfg.exists():
+        repo = ctx.obj['repo']
+        click.echo(f'Cannot find notifier configuration for Repo \'{repo.path}\'.')
+        return
+
+    if notifier_id in cfg.notifiers:
+        cfg.enable(notifier_id)
+        cfg.save()
+    else:
+        click.echo(f'No notifier with id {notifier_id} found.')
+
+
+@click.command(name='disable')
+@click.argument('notifier-id', required=True, type=str)
+@click.pass_context
+def disable_config(ctx, notifier_id):
+    """Stop receiving notifications from given notifier."""
+    cfg = ctx.obj['config']
+    if not cfg.exists():
+        repo = ctx.obj['repo']
+        click.echo(f'Cannot find notifier configuration for Repo \'{repo.path}\'.')
+        return
+
+    if notifier_id in cfg.notifiers:
+        cfg.disable(notifier_id)
+        cfg.save()
+    else:
+        click.echo(f'No notifier with id {notifier_id} found.')
+
+
+@add_config.command(name='workplace')
+@click.option('--group-id', prompt=True, required=True, type=int)
+@click.option('--access-token', prompt=True, required=True, type=str)
+@click.option('--message', prompt=MESSAGE_PROMPT, required=False, type=str,
+              default=DEFAULT_MESSAGE_TEMPLATE, show_default=True)
+@click.pass_context
+def workplace_config(ctx, group_id, access_token, message):
+    cfg = ctx.obj['config']
+    new_cfg = {
+        'id': str(uuid.uuid4()),
+        'type': 'workplace',
+        'arguments': {
+            'group_id': group_id,
+            'access_token': access_token,
+            'message': message,
+        }
+    }
+    dump_notifier_config(new_cfg)
+    click.confirm('Save notifier configuration above?', default=True, abort=True)
+    cfg.add(new_cfg)
+    cfg.save()
+
+
+@add_config.command(name='slack')
+@click.option('--webhook-url', prompt=True, required=True, type=str)
+@click.option('--message', prompt=MESSAGE_PROMPT, required=False, type=str,
+              default=DEFAULT_MESSAGE_TEMPLATE, show_default=True)
+@click.pass_context
+def slack_config(ctx, webhook_url, message):
+    cfg = ctx.obj['config']
+    new_cfg = {
+        'id': str(uuid.uuid4()),
+        'type': 'slack',
+        'arguments': {
+            'url': webhook_url,
+            'message': message,
+        }
+    }
+    dump_notifier_config(new_cfg)
+    click.confirm('Save notifier configuration above?', default=True, abort=True)
+    cfg.add(new_cfg)
+    cfg.save()
+
+
+@add_config.command(name='logger')
+@click.option('--message', prompt=MESSAGE_PROMPT, required=False, type=str,
+              default=DEFAULT_MESSAGE_TEMPLATE, show_default=True)
+@click.pass_context
+def logger_config(ctx, message):
+    cfg = ctx.obj['config']
+    new_cfg = {
+        'id': str(uuid.uuid4()),
+        'type': 'logger',
+        'arguments': {
+            'message': message,
+        }
+    }
+    dump_notifier_config(new_cfg)
+    click.confirm('Save notifier configuration above?', default=True, abort=True)
+    cfg.add(new_cfg)
+    cfg.save()
+
+
+config_notifiers.add_command(add_config)
+config_notifiers.add_command(list_config)
+config_notifiers.add_command(remove_config)
+config_notifiers.add_command(disable_config)
+config_notifiers.add_command(enable_config)
+config_notifiers.add_command(dump_config)
+config_notifiers.add_command(get_log_level)
+config_notifiers.add_command(set_log_level)

aim/ext/cleanup/__init__.py

@@ -0,0 +1,102 @@
+import time
+from abc import abstractmethod
+
+import atexit
+import weakref
+import logging
+import threading
+
+from typing import Generic, Tuple, TypeVar, Dict
+
+T = TypeVar('T')
+
+logger = logging.getLogger(__name__)
+
+
+class RobustExec(threading.Thread):
+    """
+    A thread that executes a function and catches all soft interrupts.
+    Users very often Ctrl-C to stop the program multiple times which leaves
+    program no chance to clean up.
+    """
+    def __init__(self, *args, stop_signal, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.stop_signal = stop_signal
+
+    def join(self):
+        while True:
+            try:
+                # Pressing Ctrl-C will raise KeyboardInterrupt and stop only
+                # the `Event.wait()` call. The target function of `RobustExec`
+                # will continue to run, until stop_signal is set.
+                return self.stop_signal.wait()
+            except KeyboardInterrupt:
+                logger.warning('Received Ctrl-C. Closing gracefully.')
+
+
+class AutoClean(Generic[T]):
+    PRIORITY = 10
+    _registered_with_atexit = False
+    _finalizers: Dict[
+        T, Tuple[int, weakref.finalize]
+    ] = weakref.WeakKeyDictionary()
+
+    stop_signal = threading.Event()
+
+    def __init__(self, instance: T) -> None:
+        """
+        Initialize object resources that need to be cleaned up automatically.
+        """
+        self.finalizer = weakref.finalize(instance, self._close)
+        # We do not finalizer to automatically close the object itself.
+        # Instead, we will collect all finalizers and call them at exit.
+        self.finalizer.atexit = False
+        self._finalizers[instance] = (self.PRIORITY, self.finalizer)
+
+        if not AutoClean._registered_with_atexit:
+            atexit.register(AutoClean.cleanup)
+            AutoClean._registered_with_atexit = True
+
+    def close(self) -> None:
+        """Manually close the object."""
+        self.finalizer()
+
+    @abstractmethod
+    def _close(self):
+        """
+        This method is called when the instance is garbage collected.
+        Meant to be overridden by subclasses.
+        """
+        ...
+
+    @staticmethod
+    def _cleanup():
+        """
+        Actually cleanup all the remaining resources in order of specified
+        priorities. Resources with bigger priority numbers are cleaned up first (have higher priority)
+        while the resources with equal priorities are cleaned up in reverse order
+        of their registration (hence why reversed is used).
+        This is called automatically by atexit. It is not meant tobe called manually.
+        """
+
+        finalizers = sorted(AutoClean._finalizers.items(), key=lambda x: x[1][0])
+        logger.debug(f'Cleaning up...  Found {len(finalizers)} finalizers')
+        logger.debug('Cleaning up...  Iterating over instances in order')
+        for key, (priority, finalizer) in reversed(finalizers):
+            logger.debug(f'Cleaning up...  with priority={priority} instance {key}')
+            finalizer()
+        AutoClean.stop_signal.set()
+
+    @staticmethod
+    def cleanup():
+        """
+        Cleanup all the remaining objects. This is called automatically
+        when the program exits.
+        It also blocks until all the cleanup functions have finished to ensure
+        data consistency.
+        """
+        logger.debug('Cleaning up...  Blocking KeyboardInterrupts')
+        example = RobustExec(stop_signal=AutoClean.stop_signal, target=AutoClean._cleanup)
+        example.start()
+        example.join()
+        logger.debug('Cleaning up...  Done')

aim/ext/exception_resistant.py

@@ -0,0 +1,72 @@
+import logging
+
+from functools import wraps
+
+
+logger = logging.getLogger(__name__)
+
+
+def exception_resistant(silent: bool):
+    def inner(func):
+        if not silent:
+            num_fails = 0
+            max_fails = 6
+
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                nonlocal num_fails
+                func_name = func.__name__
+                try:
+                    return func(*args, **kwargs)
+                except Exception as e:
+                    num_fails += 1
+                    if num_fails == 1:
+                        print('Something went wrong in `{}`. The process will continue to execute.'.format(func_name))
+                    if num_fails <= max_fails:
+                        print('`{}`: {}'.format(func_name, e))
+                    elif num_fails == max_fails + 1:
+                        print('The rest of the `{}` errors are hidden.'.format(func_name))
+        else:
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                try:
+                    return func(*args, **kwargs)
+                except Exception:
+                    pass
+        return wrapper
+    return inner
+
+
+class _SafeModeConfig:
+    @staticmethod
+    def log_exception(e: Exception, func: callable):
+        logger.warning(f'Exception "{str(e)}" raised in function "{func.__name__}"')
+
+    @staticmethod
+    def reraise_exception(e: Exception, func: callable):
+        raise e
+
+    exception_callback = reraise_exception
+
+
+def enable_safe_mode():
+    _SafeModeConfig.exception_callback = _SafeModeConfig.log_exception
+
+
+def disable_safe_mode():
+    _SafeModeConfig.exception_callback = _SafeModeConfig.reraise_exception
+
+
+def set_exception_callback(callback: callable):
+    _SafeModeConfig.exception_callback = callback
+
+
+def noexcept(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            _SafeModeConfig.exception_callback(e, func)
+
+    return wrapper

aim/ext/notebook/notebook.py

@@ -0,0 +1,223 @@
+import shlex
+
+from aim.cli.configs import VERSION_NAME, UP_NAME
+import aim.cli.manager.manager as manager
+
+# Error message prefix for aim commands
+ERROR_MSG_PREFIX = b'Error:'
+
+# returned by get_execution_context
+_COLAB_EXEC_CONTEXT = "_COLAB_EXEC_CONTEXT"
+_IPYTHON_EXEC_CONTEXT = "_IPYTHON_EXEC_CONTEXT"
+_OTHER_EXEC_CONTEXT = "_OTHER_EXEC_CONTEXT"
+
+# current execution context
+_CURRENT_CONTEXT = _OTHER_EXEC_CONTEXT
+
+# environment specific constants
+# useful for detecting the environment from the UI
+_SAGE_MAKER_NOTEBOOK_PATH_POSTFIX = "/aim-sage"
+_NOTEBOOK_PATH_POSTFIX = "/notebook"
+
+
+def get_execution_context():
+    """Determine the most specific context that we're in.
+    Returns:
+      _COLAB_EXEC_CONTEXT: If in Colab with an IPython notebook context.
+      _IPYTHON_EXEC_CONTEXT: If we are in an IPython notebook
+        context but not in colab (i.e. `jupyter notebook`)
+        line).
+      _OTHER_EXEC_CONTEXT: Otherwise (e.g., by running a Python script at the
+        command-line or using the `ipython` interactive shell).
+    """
+    # In Colab, the `google.colab` module is available, but the shell
+    # returned by `IPython.get_ipython` does not have a `get_trait`
+    # method.
+
+    # imports are dynamic, since some modules are not available for all contexts
+    try:
+        import IPython
+    except ImportError:
+        pass
+    else:
+        ipython = IPython.get_ipython()
+        # @TODO find a stable way to get colab context
+        if ipython is not None and 'google.colab' in str(ipython):
+            # We are in Colab notebook context
+            # global _CURRENT_CONTEXT
+            # _CURRENT_CONTEXT = _COLAB_EXEC_CONTEXT
+            return _COLAB_EXEC_CONTEXT
+
+        # In an IPython command line shell or Jupyter notebook
+        elif ipython is not None and ipython.has_trait("kernel"):
+            # global _CURRENT_CONTEXT
+            # _CURRENT_CONTEXT = _IPYTHON_EXEC_CONTEXT
+            return _IPYTHON_EXEC_CONTEXT
+
+    # Otherwise, we're not in a known notebook context.
+    return _OTHER_EXEC_CONTEXT
+
+
+def get_argument_options(line):
+    """
+        Returns parsed argument options and command from magic cell as  dict (command, options)
+        currently parse only --<name>=value style to dict
+        Set default values for the required fields, otherwise the provided fields
+        Will omit unsupported args @TODO notify about unsupported args
+        @TODO add process args all styles to dict
+    """
+    # @TODO improve this logic
+    # --proxy-url is useful to print the right url, and set UI's url into iframe correctly
+    supported_args = ['--port', '--host', '--repo', '--proxy-url']
+
+    args = shlex.split(line)
+    command = args[0]
+
+    options = {
+        '--host': '127.0.0.1',
+        '--port': '43801',
+        '--base-path': _NOTEBOOK_PATH_POSTFIX
+    }
+    for arg in args[1:]:
+        key, value = arg.split('=', 1)
+        if key in supported_args:
+            options[key] = value
+    # if --proxy-url passed
+    if options.get('--proxy-url'):
+        options['--base-path'] = f'/proxy/absolute/{options["--port"]}{_SAGE_MAKER_NOTEBOOK_PATH_POSTFIX}'
+
+    return command, options
+
+
+def display_colab(port, display):
+    """Display Aim instance in a Colab output frame.
+       It need go through the proxy
+    """
+    import IPython.display
+
+    shell = """
+        (async () => {{
+            const url = new URL('{path}/', await google.colab.kernel.proxyPort({port}, {{'cache': true}}));
+            const iframe = document.createElement('iframe');
+            iframe.src = url;
+            const a = document.createElement('a');
+            a.href = url;
+            a.innerHTML = 'Open in new browser tab';
+            a.setAttribute('target', '_blank');
+
+            iframe.setAttribute('width', '100%');
+            iframe.setAttribute('height', '800');
+            iframe.setAttribute('frameborder', 0);
+            document.body.appendChild(iframe);
+            document.body.appendChild(a);
+        }})();
+        """.format(path=_NOTEBOOK_PATH_POSTFIX, port=port)
+
+    script = IPython.display.Javascript(shell)
+
+    if display:
+        display.update(script)
+    else:
+        IPython.display.display(script)
+
+
+def display_notebook(host, port, display, proxy_url=None):
+    """Display Aim instance in an ipython context output frame.
+    """
+    import IPython.display
+    url = "{}:{}{}".format(host, port, _NOTEBOOK_PATH_POSTFIX)
+
+    # @TODO add warning if proxy_url is not defined
+    if proxy_url:
+        # jupyter-server-proxy supports absolute paths by using it with /proxy/absolute/<port> path
+        url = "{}{}{}{}/".format(proxy_url, '/proxy/absolute/', port, _SAGE_MAKER_NOTEBOOK_PATH_POSTFIX)
+        print(url)
+
+    shell = """
+              <iframe id="aim" width="100%" height="800" frameborder="0" src={}>
+              </iframe>
+            """.format(url)
+
+    iframe = IPython.display.HTML(shell)
+    display.update(iframe)
+
+
+def up(options, context):
+    """
+        Calls to run `aim up` command width corresponding arguments
+        Handles the result of the command
+        Renders the <iframe> tag for the notebook and message for the shell users
+        The <iframe> renders width the corresponding way for different execution contexts (mainly for notebooks)
+    """
+
+    try:
+        import IPython
+        import IPython.display
+    except ImportError:
+        IPython = None
+
+    display = None
+    if context == _OTHER_EXEC_CONTEXT:
+        print("Launching Aim ...")
+    else:
+        display = IPython.display.display(
+            IPython.display.Pretty("Launching Aim ..."),
+            display_id=True,
+        )
+
+    result = manager.run_process(UP_NAME, options)
+
+    if result.status == manager.ManagerActionStatuses.Failed:
+        print(result.info["message"])
+        return
+
+    port = result.info["port"]
+    host = result.info["host"]
+
+    # successful exec of aim up command
+    if context == _COLAB_EXEC_CONTEXT:
+        display_colab(port, display)
+        return
+    if context == _IPYTHON_EXEC_CONTEXT:
+        display_notebook(host, port, display, options.get("--proxy-url"))
+        return
+
+    # other context
+    print("Open {}:{}".format(host, port))
+
+
+def version(options, context):
+    """Handles aim version (get version process) and send to the ui"""
+    result = manager.run_process(VERSION_NAME, options)
+    if result.status is manager.ManagerActionStatuses.Failed:
+        print(result.info["message"])
+    else:
+        print("Aim v{}".format(result.info["version"]))
+
+
+# Those are aim magic function available commands
+# This is why we are not using constants from aim.cli.commands
+# It is possible to add commands outside aim cli
+handlers = {
+    UP_NAME: up,
+    VERSION_NAME: version
+}
+
+
+def execute_magic_aim(line):
+    """ `aim` line magic function
+        We are trying to keep similarity with the native aim cli commands as much as possible
+    """
+    context = get_execution_context()
+    command, options = get_argument_options(line)
+    # check command existence
+    if command not in handlers:
+        print('Invalid operation.')
+        return
+
+    # call corresponding handler
+    handlers[command](options, context)
+
+
+def load_ipython_extension(ipython):
+    ipython.register_magic_function(execute_magic_aim, magic_kind="line", magic_name="aim")

aim/ext/notifier/__init__.py

@@ -0,0 +1,19 @@
+from aim.ext.notifier.notifier import Notifier, NotificationSendError
+from aim.ext.notifier.notifier_builder import NotifierBuilder
+from aim.ext.notifier.config import Config
+from aim.ext.notifier.utils import get_config_path
+
+
+def get_config(base_dir) -> Config:
+    config_file = get_config_path(base_dir)
+    return Config(config_file)
+
+
+def get_notifier(base_dir) -> Notifier:
+    cfg = get_config(base_dir)
+    cfg.load()
+    builder = NotifierBuilder(cfg.notifiers)
+    return builder.build()
+
+
+__all__ = ['Notifier', 'NotificationSendError', 'Config', 'get_config', 'get_notifier']

aim/ext/notifier/base_notifier.py

@@ -0,0 +1,12 @@
+from typing import Optional
+
+
+class BaseNotifier(object):
+    def __init__(self, _id: str):
+        self._id = _id
+
+    def __repr__(self):
+        return f'<{self.__class__.__name__} object at {id(self)}>'
+
+    def notify(self, message: Optional[str] = None, **kwargs):
+        raise NotImplementedError

aim/ext/notifier/config.py

@@ -0,0 +1,69 @@
+import json
+import logging
+import uuid
+
+from typing import Dict
+from pathlib import Path
+
+from aim.ext.notifier.utils import get_empty_config_path
+
+
+class Config:
+    def __init__(self, config_file: Path):
+        self._cfg_file: Path = config_file
+        self._cfg = {}
+        if self.exists():
+            self._from_file(self._cfg_file)
+        else:
+            self._from_file(get_empty_config_path())
+
+    def exists(self) -> bool:
+        return self._cfg_file.exists()
+
+    def load(self):
+        self._from_file(self._cfg_file)
+
+    def _from_file(self, cfg_file: Path):
+        with cfg_file.open() as cfg_fh:
+            self._cfg = json.load(cfg_fh)
+
+    def save(self):
+        with self._cfg_file.open('w+') as cfg_fh:
+            json.dump(self._cfg, cfg_fh, indent=2)
+
+    def dump(self) -> str:
+        return json.dumps(self._cfg, indent=2)
+
+    def add(self, config: Dict):
+        cfg_id = config['id']
+        config['status'] = 'enabled'
+        self._cfg['notifications']['notifiers'][cfg_id] = config
+
+    def get(self, cfg_id: uuid.UUID) -> Dict:
+        return self.notifiers.get(cfg_id, {})
+
+    def enable(self, cfg_id: uuid.UUID):
+        cfg = self.notifiers[cfg_id]
+        cfg['status'] = 'enabled'
+
+    def disable(self, cfg_id: uuid.UUID):
+        cfg = self.notifiers[cfg_id]
+        cfg['status'] = 'disabled'
+
+    def remove(self, cfg_id: uuid.UUID) -> bool:
+        if cfg_id in self.notifiers:
+            del self.notifiers[cfg_id]
+            return True
+        return False
+
+    @property
+    def notifiers(self) -> Dict:
+        return self._cfg['notifications']['notifiers']
+
+    @property
+    def log_level(self) -> int:
+        return self._cfg['notifications'].get('log_level', logging.WARNING)
+
+    @log_level.setter
+    def log_level(self, lvl: int):
+        self._cfg['notifications']['log_level'] = lvl

aim/ext/notifier/config_default.json

@@ -0,0 +1,15 @@
+{
+  "notifications": {
+    "version": "1.0",
+    "notifiers": {
+    "b6487000-b521-4a2a-a42b-8fee1b7fa43d": {
+        "id": "b6487000-b521-4a2a-a42b-8fee1b7fa43d",
+        "type": "logger",
+        "status": "enabled",
+        "arguments": {
+          "message": "❗️ Something wrong with Run '{run.hash}'. Please check. ❗️"
+        }
+      }
+    }
+  }
+}

aim/ext/notifier/config_empty.json

@@ -0,0 +1,6 @@
+{
+  "notifications": {
+    "version": "1.0",
+    "notifiers": {}
+  }
+}

aim/ext/notifier/logging_notifier.py

@@ -0,0 +1,16 @@
+import logging
+from typing import Optional
+
+from aim.ext.notifier.base_notifier import BaseNotifier
+
+
+class LoggingNotifier(BaseNotifier):
+    def __init__(self, _id: str, config: dict):
+        super().__init__(_id)
+        self.message_template = config['message']
+        self.logger = logging.getLogger('notifier')
+
+    def notify(self, message: Optional[str] = None, **kwargs):
+        message_template = message or self.message_template
+        msg = message_template.format(**kwargs)
+        self.logger.error(msg)

aim/ext/notifier/notifier.py

@@ -0,0 +1,42 @@
+import logging
+import time
+from typing import List, Optional
+
+from aim.ext.notifier.base_notifier import BaseNotifier
+
+logger = logging.getLogger(__name__)
+
+
+class NotificationSendError(RuntimeError):
+    def __init__(self, e: Exception):
+        super().__init__(e)
+
+
+class Notifier(BaseNotifier):
+    MAX_RETRIES = 5
+    RETRY_DELAY = 1
+
+    def __init__(self):
+        self._notifiers: List[BaseNotifier] = []
+
+    def add(self, sub: BaseNotifier):
+        self._notifiers.append(sub)
+
+    def notify(self, message: Optional[str] = None, **kwargs):
+        for sub in self._notifiers:
+            attempt = 0
+            while attempt < self.MAX_RETRIES:
+                try:
+                    sub.notify(message, **kwargs)
+                    break
+                except Exception as e:
+                    attempt += 1
+                    if attempt == self.MAX_RETRIES:
+                        logger.error(f'Notifier {sub} failed to send message "{message}". '
+                                     f'No retries left.')
+                        raise NotificationSendError(e)
+                    else:
+                        logger.error(f'Notifier {sub} failed to send message "{message}". '
+                                     f'Retry attempts left {self.MAX_RETRIES - attempt} '
+                                     f'Next retry in {self.RETRY_DELAY} seconds.')
+                        time.sleep(self.RETRY_DELAY)

aim/ext/notifier/notifier_builder.py

@@ -0,0 +1,36 @@
+import logging
+
+from aim.ext.notifier.logging_notifier import LoggingNotifier
+from aim.ext.notifier.slack_notifier import SlackNotifier
+from aim.ext.notifier.workplace_notifier import WorkplaceNotifier
+from aim.ext.notifier.notifier import Notifier
+
+logger = logging.getLogger(__name__)
+
+
+class NotifierBuilder(object):
+    _factories = {
+        'logger': LoggingNotifier,
+        'workplace': WorkplaceNotifier,
+        'slack': SlackNotifier
+    }
+
+    def __init__(self, config: dict):
+        self.config = config
+
+    def build(self) -> Notifier:
+        notifier = Notifier()
+        for sub_config in self.config.values():
+            notif_type = sub_config['type']
+            notif_id = sub_config['id']
+            args = sub_config['arguments']
+            if sub_config['status'] != 'enabled':
+                continue
+            try:
+                notif_cls = self._factories[notif_type]
+            except KeyError:
+                logger.warning(f'Unknown notifier type {notif_type}. Skipping.')
+            else:
+                sub_notifier = notif_cls(notif_id, args)
+                notifier.add(sub_notifier)
+        return notifier

aim/ext/notifier/slack_notifier.py

@@ -0,0 +1,17 @@
+import requests
+from typing import Optional
+
+from aim.ext.notifier.base_notifier import BaseNotifier
+
+
+class SlackNotifier(BaseNotifier):
+    def __init__(self, _id: str, config: dict):
+        super().__init__(_id)
+        self.message_template = config['message']
+
+        self.url = config['url']
+
+    def notify(self, message: Optional[str] = None, **kwargs):
+        message_template = message or self.message_template
+        msg = message_template.format(**kwargs)
+        requests.post(self.url, json={'text': msg})

aim/ext/notifier/utils.py

@@ -0,0 +1,40 @@
+import os
+import shutil
+
+from functools import lru_cache
+from pathlib import Path
+from typing import Union
+
+here = os.path.abspath(os.path.dirname(__file__))
+
+
+@lru_cache()
+def get_working_directory(base_dir: Union[str, Path]) -> Path:
+    if not isinstance(base_dir, Path):
+        base_dir = Path(base_dir)
+    work_dir = base_dir / 'ext' / 'notifications'
+    work_dir.mkdir(parents=True, exist_ok=True)
+    return work_dir
+
+
+def get_config_path(base_dir: Union[str, Path]) -> Path:
+    return get_working_directory(base_dir) / 'config.json'
+
+
+def get_default_config_path() -> Path:
+    return Path(here) / 'config_default.json'
+
+
+def get_empty_config_path() -> Path:
+    return Path(here) / 'config_empty.json'
+
+
+def has_watcher_config(base_dir: Union[str, Path]) -> bool:
+    config_path = get_config_path(base_dir)
+    return config_path.is_file()
+
+
+def set_default_config(base_dir: Union[str, Path]):
+    config = get_config_path(base_dir)
+    default_config = get_default_config_path()
+    shutil.copy(default_config, config)

aim/ext/notifier/workplace_notifier.py

@@ -0,0 +1,24 @@
+import requests
+from typing import Optional
+
+from aim.ext.notifier.base_notifier import BaseNotifier
+
+
+class WorkplaceNotifier(BaseNotifier):
+    def __init__(self, _id: str, config: dict):
+        super().__init__(_id)
+        self.message_template = config['message']
+        self.wp_access_token = config['access_token']
+        self.wp_group_id = config['group_id']
+        self.url = self._get_workplace_url()
+
+    def notify(self, message: Optional[str] = None, **kwargs):
+        message_template = message or self.message_template
+        msg = message_template.format(**kwargs)
+        params = {'access_token': self.wp_access_token}
+        requests.post(self.url, json={'message': msg}, params=params)
+
+    def _get_workplace_url(self):
+        api_version = 'v14.0'
+
+        return f'https://graph.facebook.com/{api_version}/{self.wp_group_id}/feed'

aim/ext/resource/__init__.py

@@ -0,0 +1,2 @@
+from aim.ext.resource.tracker import ResourceTracker
+from aim.ext.resource.configs import *

aim/ext/resource/configs.py

@@ -0,0 +1,3 @@
+# Aim Resource Collector
+AIM_RESOURCE_METRIC_PREFIX = '__system__'
+DEFAULT_SYSTEM_TRACKING_INT = 10

aim/ext/resource/log.py

@@ -0,0 +1,29 @@
+from typing import Tuple, Union
+
+from aim.storage.object import CustomObject
+from aim.sdk.sequence import MediaSequenceBase
+from aim.sdk.objects.text import Text
+
+
+@CustomObject.alias('aim.log_line')
+class LogLine(Text):
+    AIM_NAME = 'aim.log_line'
+
+    def __init__(self, line):
+        if not isinstance(line, str):
+            raise TypeError('`line` must be a string.')
+
+        super().__init__(line)
+
+
+class Logs(MediaSequenceBase):
+    """Class representing series of LogLine objects."""
+
+    @classmethod
+    def allowed_dtypes(cls) -> Union[str, Tuple[str, ...]]:
+        text_typename = LogLine.get_typename()
+        return text_typename
+
+    @classmethod
+    def sequence_name(cls) -> str:
+        return 'logs'

aim/ext/resource/stat.py

@@ -0,0 +1,191 @@
+import psutil
+import json
+from typing import List
+
+from aim.ext.resource.utils import round10e5
+
+import aim.ext.pynvml as nvml
+
+
+class StatDict(object):
+    # Available aggregation functions
+    AGG_MODE_AVG = 'average'
+    AGG_MODE_MIN = 'min'
+    AGG_MODE_MAX = 'max'
+    AGG_MODE_DIFF = 'diff'
+    AGG_DEFAULT = AGG_MODE_AVG
+
+    @classmethod
+    def aggregate(cls, items: List, mode: str):
+        """
+        Aggregates array of numbers by a given 'mode'
+        """
+        if mode == cls.AGG_MODE_MAX:
+            return max(items)
+        elif mode == cls.AGG_MODE_MIN:
+            return min(items)
+        elif mode == cls.AGG_MODE_AVG:
+            return round10e5(sum(items) / len(items))
+        elif mode == cls.AGG_MODE_DIFF:
+            return round10e5(max(items) - min(items))
+        else:
+            raise ValueError('unknown aggregation mode: \'{}\''.format(mode))
+
+    @classmethod
+    def aggregate_items(cls,
+                        items: 'List[StatDict]',
+                        agg_mode: str = AGG_DEFAULT,
+                        ):
+        """
+        Aggregates array of `StatDict` items by a given `mode`
+        """
+        aggregated_stat = cls()
+
+        # Return empty item if items array is empty
+        if not items or len(items) == 0:
+            return aggregated_stat
+
+        gpu_stats = []
+        for s in items:
+            # Collect system stats
+            for k in s.system.keys():
+                aggregated_stat.system.setdefault(k, [])
+                aggregated_stat.system[k].append(s.system[k])
+
+            # Collect GPU device stats
+            for stat_item_gpu_idx in range(len(s.gpus)):
+                stat_item_gpu_stat = s.gpus[stat_item_gpu_idx]
+                if len(gpu_stats) == stat_item_gpu_idx:
+                    gpu_stats.append({})
+                for gpu_stat_key in stat_item_gpu_stat.keys():
+                    gpu_stat = stat_item_gpu_stat[gpu_stat_key]
+                    gpu_stats[stat_item_gpu_idx].setdefault(gpu_stat_key, [])
+                    gpu_stats[stat_item_gpu_idx][gpu_stat_key].append(gpu_stat)
+
+        # Aggregate system stats
+        for k in aggregated_stat.system.keys():
+            aggregated_stat.system[k] = cls.aggregate(aggregated_stat.system[k],
+                                                      agg_mode)
+
+        # Aggregate GPU device stats
+        for g in range(len(gpu_stats)):
+            for k in gpu_stats[g].keys():
+                gpu_stats[g][k] = cls.aggregate(gpu_stats[g][k], agg_mode)
+        aggregated_stat.gpu = gpu_stats
+
+        return aggregated_stat
+
+    def __init__(self, system: dict = None, gpus: List[dict] = None):
+        self.system = system or {}
+        self.gpus = gpus or []
+
+    def __str__(self):
+        return json.dumps(self.to_dict())
+
+    def to_dict(self):
+        """
+        Returns system and GPU device statistics
+        """
+        return {
+            'system': self.system,
+            'gpus': self.gpus,
+        }
+
+
+class Stat(object):
+    def __init__(self, process):
+        # Set process
+        self._process = process
+
+        # Get statistics
+        system, gpus = self.get_stats()
+        self._stat = StatDict(system, gpus)
+
+    @property
+    def process(self):
+        return self._process
+
+    @property
+    def stat_item(self):
+        return self._stat
+
+    @property
+    def system(self):
+        return self._stat.system
+
+    @property
+    def gpus(self):
+        return self._stat.gpus
+
+    def get_stats(self):
+        """
+        Get system statistics and assign to `self`
+        """
+        memory_usage = psutil.virtual_memory()
+        disk_usage = psutil.disk_usage('/')
+        # net = psutil.net_io_counters()
+        system = {
+            # CPU utilization percent(can be over 100%)
+            'cpu': round10e5(self._process.cpu_percent(0.0)),
+
+            # Whole system memory usage
+            # 'memory_used': round10e5(memory_usage.used / 1024 / 1024),
+            'memory_percent': round10e5(memory_usage.used * 100 / memory_usage.total),
+
+            # Get the portion of memory occupied by a process
+            # 'p_memory_rss': round10e5(self._process.memory_info().rss
+            #                           / 1024 / 1024),
+            'p_memory_percent': round10e5(self._process.memory_percent()),
+
+            # Disk usage
+            # 'disk_used': round10e5(disk_usage.used / 1024 / 1024),
+            'disk_percent': round10e5(disk_usage.percent),
+        }
+
+        # Collect GPU statistics
+        gpus = []
+        try:
+            nvml.nvmlInit()
+            gpu_device_count = nvml.nvmlDeviceGetCount()
+            for i in range(gpu_device_count):
+                gpu_info = dict()
+                handle = nvml.nvmlDeviceGetHandleByIndex(i)
+                try:
+                    util = nvml.nvmlDeviceGetUtilizationRates(handle)
+                    # GPU utilization percent
+                    gpu_info["gpu"] = round10e5(util.gpu)
+                except nvml.NVMLError_NotSupported:
+                    pass
+                try:
+                    # Get device memory
+                    memory = nvml.nvmlDeviceGetMemoryInfo(handle)
+                    # Device memory usage
+                    # 'memory_used': round10e5(memory.used / 1024 / 1024),
+                    gpu_info["gpu_memory_percent"] = round10e5(memory.used * 100 / memory.total)
+                except nvml.NVMLError_NotSupported:
+                    pass
+                try:
+                    # Get device temperature
+                    nvml_tmp = nvml.NVML_TEMPERATURE_GPU
+                    temp = nvml.nvmlDeviceGetTemperature(handle, nvml_tmp)
+                    # Device temperature
+                    gpu_info["gpu_temp"] = round10e5(temp)
+                except nvml.NVMLError_NotSupported:
+                    pass
+                try:
+                    # Compute power usage in watts and percent
+                    power_watts = nvml.nvmlDeviceGetPowerUsage(handle) / 1000
+                    power_cap = nvml.nvmlDeviceGetEnforcedPowerLimit(handle)
+                    power_cap_watts = power_cap / 1000
+                    power_watts / power_cap_watts * 100
+                    # Power usage in watts and percent
+                    gpu_info["gpu_power_watts"] = round10e5(power_watts)
+                    # gpu_info["power_percent"] = round10e5(power_usage)
+                except nvml.NVMLError_NotSupported:
+                    pass
+                gpus.append(gpu_info)
+            nvml.nvmlShutdown()
+        except nvml.NVMLError:
+            pass
+
+        return system, gpus

aim/ext/resource/tracker.py

@@ -0,0 +1,233 @@
+import io
+import logging
+import re
+import sys
+import time
+import weakref
+
+from psutil import Process, cpu_percent
+from threading import Thread
+from typing import Union
+from weakref import WeakValueDictionary
+
+from aim.ext.resource.configs import AIM_RESOURCE_METRIC_PREFIX
+from aim.ext.resource.log import LogLine
+from aim.ext.resource.stat import Stat
+
+logger = logging.getLogger(__name__)
+
+
+class ResourceTracker(object):
+    _buffer_registry = WeakValueDictionary()
+    _old_out_write = None
+    _old_err_write = None
+
+    @classmethod
+    def _install_stream_patches(cls):
+        cls._old_out_write = sys.stdout.write
+        cls._old_err_write = sys.stderr.write
+
+        def new_out_write(data):
+            cls._old_out_write(data)
+            if isinstance(data, str):
+                data = data.encode()
+            for buffer in cls._buffer_registry.values():
+                buffer.write(data)
+
+        def new_err_write(data):
+            cls._old_err_write(data)
+            if isinstance(data, str):
+                data = data.encode()
+            for buffer in cls._buffer_registry.values():
+                buffer.write(data)
+
+        sys.stdout.write = new_out_write
+        sys.stderr.write = new_err_write
+
+    @classmethod
+    def _uninstall_stream_patches(cls):
+        sys.stdout.write = cls._old_out_write
+        sys.stderr.write = cls._old_err_write
+
+    STAT_INTERVAL_MIN = 0.1
+    STAT_INTERVAL_MAX = 24 * 60 * 60.0
+    STAT_INTERVAL_DEFAULT = 60.0
+
+    @classmethod
+    def check_interval(cls, interval, warn=True):
+        if interval is None:
+            warn = False
+        if not isinstance(interval, (int, float)) or not cls.STAT_INTERVAL_MIN <= interval <= cls.STAT_INTERVAL_MAX:
+            if warn:
+                logger.warning('To track system resource usage '
+                               'please set `system_tracking_interval` '
+                               'greater than 0 and less than 1 day')
+            return False
+        return True
+
+    reset_cpu_cycle = False
+
+    @staticmethod
+    def reset_proc_interval():
+        """
+        Calls process `cpu_percent` which resets cpu utilization tracking cycle
+        Read more: https://psutil.readthedocs.io/en/latest/#psutil.cpu_percent
+        """
+        cpu_percent(0.0)
+
+    def __init__(self,
+                 tracker,
+                 interval: Union[int, float] = STAT_INTERVAL_DEFAULT,
+                 capture_logs: bool = True,
+                 log_offset: int = 0):
+        self._tracker = weakref.ref(tracker)
+        self._stat_capture_interval = None
+        if self.check_interval(interval, warn=False):
+            self._stat_capture_interval = interval
+
+        # terminal log capturing
+        self._capture_logs = capture_logs
+        self._log_capture_interval = 1
+        self._old_out = None
+        self._old_err = None
+        self._io_buffer = io.BytesIO()
+        self._line_counter = log_offset
+
+        try:
+            self._process = Process()
+        except Exception:
+            self._process = None
+
+        # Start thread to collect stats and logs at intervals
+        self._th_collector = Thread(target=self._stat_collector, daemon=True)
+        self._shutdown = False
+        self._started = False
+
+        if ResourceTracker.reset_cpu_cycle is False:
+            ResourceTracker.reset_cpu_cycle = True
+            self.reset_proc_interval()
+
+    def start(self):
+        """
+        Start statistics collection
+        """
+        if self._started:
+            return
+
+        self._started = True
+        if self._capture_logs:
+            # install the stream patches if not done yet
+            if not self._buffer_registry:
+                self._install_stream_patches()
+            self._buffer_registry[id(self)] = self._io_buffer
+        # Start thread to asynchronously collect statistics
+        self._th_collector.start()
+
+    def stop(self):
+        if not self._started:
+            return
+
+        self._shutdown = True
+        self._th_collector.join()
+        if self._capture_logs:
+            # read and store remaining buffered logs
+            self._store_buffered_logs()
+            # unregister the buffer
+            del self._buffer_registry[id(self)]
+            # uninstall stream patching if no buffer is left in the registry
+            if not self._buffer_registry:
+                self._uninstall_stream_patches()
+
+    def close(self):
+        """Interface to make compatible with Resource AutoClean"""
+        self.stop()
+
+    def _track(self, stat: Stat):
+        # Store system stats
+        for resource, usage in stat.system.items():
+            self._tracker()(
+                usage,
+                name='{}{}'.format(AIM_RESOURCE_METRIC_PREFIX, resource),
+            )
+
+        # Store GPU stats
+        for gpu_idx, gpu in enumerate(stat.gpus):
+            for resource, usage in gpu.items():
+                self._tracker()(
+                    usage,
+                    name='{}{}'.format(AIM_RESOURCE_METRIC_PREFIX, resource),
+                    context={'gpu': gpu_idx}
+                )
+
+    def _stat_collector(self):
+        """
+            Statistics collecting thread body
+        """
+        stat_time_counter = 0
+        log_capture_time_counter = 0
+
+        # store initial system usage stats
+        if self._stat_capture_interval:
+            stat = Stat(self._process)
+            self._track(stat)
+
+        while True:
+            # Get system statistics
+            if self._shutdown:
+                break
+
+            time.sleep(0.1)
+            stat_time_counter += 0.1
+            log_capture_time_counter += 0.1
+
+            if self._stat_capture_interval and stat_time_counter > self._stat_capture_interval:
+                stat = Stat(self._process)
+                self._track(stat)
+                stat_time_counter = 0
+
+            if self._capture_logs and log_capture_time_counter > self._log_capture_interval:
+                self._store_buffered_logs()
+                log_capture_time_counter = 0
+
+    def _store_buffered_logs(self):
+        _buffer_size = self._io_buffer.tell()
+        if not _buffer_size:
+            return
+
+        self._io_buffer.seek(0)
+        # read and reset the buffer
+        data = self._io_buffer.read(_buffer_size)
+        self._io_buffer.seek(0)
+        # handle the buffered data and store
+
+        lines = data.split(b'\n')
+        ansi_csi_re = re.compile(b"\001?\033\\[((?:\\d|;)*)([a-dA-D])\002?")
+
+        def _handle_csi(line):
+            def _remove_csi(line):
+                return re.sub(ansi_csi_re, b'', line)
+
+            for match in ansi_csi_re.finditer(line):
+                arg, command = match.groups()
+                arg = int(arg.decode()) if arg else 1
+                if command == b'A':  # cursor up
+                    self._line_counter -= arg
+                if command == b'B':  # cursor down
+                    self._line_counter += arg
+
+            return _remove_csi(line)
+
+        line = None
+        for line in lines:
+            # handle cursor up and down symbols
+            line = _handle_csi(line)
+            # handle each line for carriage returns
+            line = line.rsplit(b'\r')[-1]
+            self._tracker()(LogLine(line.decode()), name='logs', step=self._line_counter)
+            self._line_counter += 1
+
+        self._line_counter -= 1
+
+        # if there was no b'\n' at the end of the data keep the last line in buffer for further writing
+        if line != b'':
+            self._io_buffer.write(line)

aim/ext/resource/utils.py

@@ -0,0 +1,2 @@
+def round10e5(val):
+    return round(val * 10e5) / 10e5

aim/ext/sshfs/utils.py

@@ -0,0 +1,196 @@
+import os
+import platform
+import shutil
+import subprocess
+import time
+import logging
+
+from typing import Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+def get_username():
+    """
+    Utility function to retrieve the current username
+    """
+    if os.environ.get('LOGNAME'):
+        return os.environ.get('LOGNAME')
+    if os.environ.get('LNAME'):
+        return os.environ.get('LNAME')
+    if os.environ.get('USER'):
+        return os.environ.get('USER')
+    if os.environ.get('USERNAME'):
+        return os.environ.get('USERNAME')
+
+    return None
+
+
+def check_sshfs_installation():
+    """
+    Utility function to check sshfs installation and user permission to execute sshfs.
+    Raises errors with hints to install sshfs based on the platform
+    """
+    sshfs_executable = shutil.which('sshfs')
+    if not sshfs_executable:
+        if platform.system().lower() == 'darwin':
+            raise subprocess.SubprocessError('Could not find sshfs installation. \n'
+                                             'To install sshfs please run the following commands:\n'
+                                             'brew install macfuse \n'
+                                             'brew install gromgit/fuse/sshfs\n'
+                                             'or alternatively follow the instructions here:'
+                                             'https://osxfuse.github.io/')
+        else:
+            raise subprocess.SubprocessError('Could not find sshfs installation. \n'
+                                             'To install sshfs please run the following commands:\n'
+                                             'apt-get update\n'
+                                             'apt-get install sshfs')
+
+    # check user permission to execute sshfs
+    if not os.access(sshfs_executable, os.EX_OK):
+        username = get_username()
+        raise PermissionError(f'User {username} does not have sufficient permissions to run sshfs command.')
+
+
+def check_directory_permissions(path: str, mount_root: Optional[str] = '', unmount: bool = False):
+    """
+    Utility function to check user write permissions on given path.
+    Optionally can unmount the the given path.
+    """
+    if not os.access(path, os.W_OK):
+        username = get_username()
+        if unmount:
+            unmount_remote_repo(path, mount_root)
+        raise PermissionError(f'User {username} does not have write permissions on path: {path}.')
+
+
+def extract_remote_user_pwd_host_repo(path: str) -> Tuple[Optional[str], Optional[str], Optional[str], str]:
+    """
+    Utility function to extract user, password, host and remote repo path from given remote path.
+    Arg path should follow the following format: 'ssh://[user:password@][host]:remote_repo_path'
+    Falls back to environment variables if some of the extractables are missing
+    """
+    assert path.startswith('ssh://')
+    # strip the 'ssh://' prefix
+    path = path[6:]
+
+    # split user:password and host:remote_repo_path pairs
+    initial_split = path.split('@', maxsplit=1)
+
+    user = None
+    password = None
+    host = None
+
+    if len(initial_split) == 1:
+        user_pwd_split = None
+        host_path_split = initial_split[0].split(':', maxsplit=1)
+    else:
+        host_path_split = initial_split[1].split(':', maxsplit=1)
+        user_pwd_split = initial_split[0].split(':', maxsplit=1)
+
+    # split user and password
+    if user_pwd_split:
+        user = user_pwd_split[0]
+        if len(user_pwd_split) > 1:
+            password = user_pwd_split[1]
+
+    # split host and remote path
+    if len(host_path_split) == 1:
+        remote_path = host_path_split[0]
+    else:
+        host = host_path_split[0]
+        remote_path = host_path_split[1]
+
+    # fallback to env variables
+    if not user and os.environ.get('AIM_REMOTE_REPO_USER'):
+        user = os.environ.get('AIM_REMOTE_REPO_USER')
+    if not password and os.environ.get('AIM_REMOTE_REPO_PASSWORD'):
+        password = os.environ.get('AIM_REMOTE_REPO_PASSWORD')
+    if not host and os.environ.get('AIM_REMOTE_REPO_HOST'):
+        host = os.environ.get('AIM_REMOTE_REPO_HOST')
+
+    return user, password, host, remote_path
+
+
+def mount_remote_repo(remote_path: str) -> Tuple[str, str]:
+    """
+    Utility function to mount remote repository path to local  '/tmp/<timestamp>/<remote_repo_name>' directory
+    """
+
+    check_sshfs_installation()
+
+    user, pwd, host, remote_repo_path = extract_remote_user_pwd_host_repo(remote_path)
+
+    # check if the current user has permissions to perform write operations on /tmp dir
+    check_directory_permissions(path='/tmp')
+    mount_root = f'/tmp/{int(time.time())}'
+    mount_point = f'{mount_root}/{remote_repo_path}'
+    os.makedirs(mount_point, exist_ok=True)
+
+    login_options = f'{host}:{remote_repo_path}'
+    if user:
+        if pwd:
+            login_options = f'{user}:{pwd}@{login_options}'
+        else:
+            login_options = f'{user}@{login_options}'
+
+    # TODO: experiment on sshfs options to find an optimal set
+    # construct sshfs options, use IdentityFile if provided by the user
+    sshfs_options = 'default_permissions'
+    identity_file = os.environ.get('AIM_REMOTE_REPO_KEY_FILE')
+    if identity_file:
+        sshfs_options = f'{sshfs_options},IdentityFile={identity_file}'
+
+    # try mounting using sshfs
+    cmd = ['sshfs', login_options, mount_point, '-o', sshfs_options]
+    sshfs_process = subprocess.Popen(cmd, stderr=subprocess.PIPE)
+    # redirect sdterr to pipe and check if errors occurred, if not let the sshfs process run in the background
+    # otherwise raise an error, kill the process and cleanup the local mount point
+    if sshfs_process.stderr.readline():
+        # if mounting fails remove local mount point created by Aim
+        # and print the command that was used, so it'll be easier for the user to perform manual mounting
+        sshfs_process.wait()
+        unmount_remote_repo(mount_point, mount_root)
+        raise subprocess.SubprocessError(f'Could not mount remote repository using command: \n'
+                                         f'{" ".join(cmd)} \n'
+                                         f'Please try to mount manually '
+                                         f'and use the local mount point as a usual repo path.')
+
+    # check if the user has permissions to perform write operations on local mount point (which means that the
+    # remote user also can perform write operations on remote repo path as well)
+    # unmount the local mount point in case of failure
+    check_directory_permissions(mount_point, mount_root, unmount=True)
+
+    return mount_root, mount_point
+
+
+def unmount_remote_repo(mount_point: str, mount_root: str):
+    """
+    Utility function to unmount remote repo and cleanup local directories created by Aim for mounting
+    """
+
+    # TODO: [MV] this is experimental
+    # TODO: [MV] decide later what to do with unmounting remote repo
+    if not os.path.exists(mount_point):
+        return
+
+    if not os.path.ismount(mount_point):
+        shutil.rmtree(mount_root)
+        return
+
+    # force unmount the remote path as we are sure we've done everything we need up until to this point
+    # and don't want to get stopped by any other processes stopping from unmounting
+    umount_executable = shutil.which('umount')
+    cmd = [umount_executable, mount_point]
+    child = subprocess.Popen(
+        cmd,
+    )
+    child.communicate()
+    exit_code = child.wait()
+    if exit_code != 0:
+        # in case of failure log warning so the user can unmount manually if needed
+        logger.warning(f'Could not unmount path: {mount_point}.\n'
+                       f'Please unmount manually using command:\n'
+                       f'{" ".join(cmd)}')
+    else:
+        shutil.rmtree(mount_root)

aim/ext/task_queue/queue.py

@@ -0,0 +1,64 @@
+import threading
+import queue
+import atexit
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class TaskQueue(object):
+    def __init__(self, name, num_workers=1, max_backlog=0):
+        self.name = name
+        self.max_backlog = max_backlog
+        self.num_workers = num_workers
+
+        self._queue = queue.Queue(maxsize=max_backlog)
+
+        self._threads = []
+        self._shutdown = False
+        self._stopped = False
+        atexit.register(self.stop_workers)
+
+        for thread_num in range(self.num_workers):
+            thread = threading.Thread(target=self.worker)
+            thread.daemon = True
+            self._threads.append(thread)
+            thread.start()
+
+    def register_task(self, task_func, *args, **kwargs):
+        warn_queue_full = False
+        if self._stopped:
+            logger.debug('Cannot register task: task queue is stopped.')
+        else:
+            backlog_size = self._queue.qsize()
+            if backlog_size > self.max_backlog * 0.8:  # queue is 80% full
+                warn_queue_full = True
+            self._queue.put((task_func, args, kwargs))
+        return warn_queue_full
+
+    def worker(self):
+        while True:
+            if self._shutdown:
+                logger.debug(f'Shutting down worker thread {threading.get_ident()}.')
+                break
+            task_f, args, kwargs = self._queue.get()
+            task_f(*args, **kwargs)
+            # clear the unnecessary references to Run objects
+            task_f, args, kwargs = None, None, None
+            self._queue.task_done()
+
+    def stop_workers(self):
+        if self._stopped:
+            return
+
+        self._stopped = True
+        pending_task_count = self._queue.qsize()
+        if pending_task_count:
+            logger.warning(f'Processing {pending_task_count} pending tasks in queue \'{self.name}\'... '
+                           f'Please do not kill the process.')
+            self._queue.join()
+        self._shutdown = True
+        logger.debug('No pending tasks left.')
+
+    def __del__(self):
+        self.stop_workers()

aim/ext/tensorboard_tracker/__init__.py

@@ -0,0 +1 @@
+from aim.ext.tensorboard_tracker.run import Run

aim/ext/tensorboard_tracker/run.py

@@ -0,0 +1,32 @@
+from typing import Optional, Union
+
+from aim.sdk.run import Run as SdkRun
+from aim.ext.tensorboard_tracker.tracker import TensorboardTracker
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from aim.sdk.repo import Repo
+
+
+class Run(SdkRun):
+    def __init__(
+            self, run_hash: Optional[str] = None, *,
+            sync_tensorboard_log_dir: str,
+            repo: Optional[Union[str, 'Repo']] = None,
+            experiment: Optional[str] = None,
+            force_resume: Optional[bool] = False,
+            system_tracking_interval: Optional[Union[int, float]] = None,
+            log_system_params: Optional[bool] = False,
+            capture_terminal_logs: Optional[bool] = False,
+    ):
+        super().__init__(
+            run_hash, repo=repo, read_only=False, experiment=experiment, force_resume=force_resume,
+            system_tracking_interval=system_tracking_interval, log_system_params=log_system_params,
+            capture_terminal_logs=capture_terminal_logs
+        )
+
+        self['tb_log_directory'] = sync_tensorboard_log_dir
+        self._tensorboard_tracker = TensorboardTracker(self._tracker, sync_tensorboard_log_dir)
+        self._tensorboard_tracker.start()
+        self._resources.add_extra_resource(self._tensorboard_tracker)

aim/ext/tensorboard_tracker/tracker.py

@@ -0,0 +1,255 @@
+from tensorboard.backend.event_processing.directory_watcher import DirectoryWatcher
+from tensorboard.backend.event_processing import event_file_loader
+import tensorflow as tf
+from tensorboard.util import tensor_util
+import time
+import threading
+from pathlib import Path
+import logging
+import os
+import weakref
+import queue
+
+from typing import Any
+from aim import Audio, Image, Distribution
+
+
+def _decode_histogram(value):
+    """
+    From the tensorflow histogram representation (not plugin), create an aim Distribution
+
+    :param value: value with `histo` property
+    :return: aim Distribution
+    """
+    bin_counts = list(value.histo.bucket)
+    bucket_limits = list(value.histo.bucket_limit)
+
+    if (len(bin_counts) <= 2) or (len(bucket_limits) < 2) or (bucket_limits[0] == bucket_limits[-1]):
+        return None
+
+    # This is a bit weird but it seems the histogram counts is usually padded by 0 as tensorboard
+    # only stores the right limits?
+    # See https://github.com/pytorch/pytorch/blob/7d2a18da0b3427fcbe44b461a0aa508194535885/torch/utils/tensorboard/summary.py#L390 # noqa
+    bin_counts = bin_counts[1:]
+
+    bin_range = (bucket_limits[0], bucket_limits[-1])
+    track_val = Distribution(hist=bin_counts, bin_range=bin_range)
+    return track_val
+
+
+def _decode_histogram_from_plugin(value):
+    """
+    Convert from tensorflow histogram plugin representation of the data as a tensor back into
+    a `aim` `Distribution`
+
+    Representation of histogram given by tf summary is obtained from here:
+    https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/histogram/summary_v2.py
+
+    :param value: value with a tensor that contains three columns, left_edge, right_edge,
+                  bin_values
+    :return: aim Distribution
+    """
+    left_right_bins = tensor_util.make_ndarray(value.tensor)
+    if left_right_bins is None:
+        return None
+
+    left_edge = left_right_bins[:, 0]
+    right_edge = left_right_bins[:, 1]
+    bin_counts = left_right_bins[:, 2]
+
+    bin_range = (left_edge[0], right_edge[-1])
+
+    is_empty = False
+    is_empty |= (left_right_bins.shape[0] == 0)
+    is_empty |= (bin_range[0] == bin_range[1])
+    if is_empty:
+        return None
+
+    track_val = Distribution(hist=bin_counts, bin_range=bin_range)
+    return track_val
+
+
+class TensorboardTracker:
+    def __init__(self, tracker, sync_tensorboard_log_dir: str) -> None:
+        self.tracker = tracker
+        self.sync_tensorboard_log_dir = sync_tensorboard_log_dir
+        self.tensorboard_folder_watchers = []
+        self._thread = threading.Thread(target=self._monitor_eventfiles, daemon=True)
+        self.directories_track_status = {}
+        self._shutdown = False
+        self._started = False
+        self._watcher_queue = queue.Queue()
+
+    def _monitor_eventfiles(self):
+        while True:
+            if self._shutdown:
+                break
+            for event_file in set(Path(self.sync_tensorboard_log_dir).rglob("*.tfevents*")):
+                dir = str(event_file.parent.absolute())
+                if dir not in self.directories_track_status:
+                    self.directories_track_status[dir] = "NOT_STARTED"
+            for dir, status in self.directories_track_status.items():
+                if status == "NOT_STARTED":
+                    tensorboard_folder_watcher = TensorboardFolderTracker(dir, self._watcher_queue)
+                    tensorboard_folder_watcher.start()
+                    self.tensorboard_folder_watchers.append(tensorboard_folder_watcher)
+                    self.directories_track_status[dir] = "STARTED"
+            time.sleep(5)
+
+    def start(self):
+        if self._started:
+            return
+        self._started = True
+        self._thread.start()
+        self._consumer = TensorboardEventConsumer(
+            self._watcher_queue, self.tracker
+        )
+        self._consumer.start()
+
+    def stop(self):
+        if not self._started:
+            return
+        self._shutdown = True
+        self._thread.join()
+        for tensorboard_folder_watcher in self.tensorboard_folder_watchers:
+            tensorboard_folder_watcher.stop()
+        self._consumer.stop()
+
+    def close(self):
+        """Interface to make compatible with Resource AutoClean"""
+        self.stop()
+
+
+class TensorboardFolderTracker:
+    def __init__(self, tensorboard_event_folder: str, queue: queue.Queue) -> None:
+        self.queue = queue
+        self.supported_plugins = ("images", "scalars", "histograms")
+        self.unsupported_plugin_noticed = False
+        self.folder_name = os.path.basename(tensorboard_event_folder)
+        self._thread = threading.Thread(target=self._process_event)
+        self._generator = DirectoryWatcher(tensorboard_event_folder, event_file_loader.EventFileLoader)
+        self._shutdown = False
+        self._started = False
+
+    def start(self):
+        if self._started:
+            return
+        self._started = True
+        self._thread.start()
+
+    def stop(self):
+        if not self._started:
+            return
+        self._shutdown = True
+        self._thread.join()
+
+    def _process_event(self):
+        while True:
+            if self._shutdown:
+                break
+            for event in self._generator.Load():
+                self._process_tb_event(event)
+            time.sleep(1)
+
+    def _process_tb_event(self, event):
+        def create_ndarray(tensor):
+            res = tensor_util.make_ndarray(tensor)
+            if res.dtype == "object":
+                return None
+            else:
+                return res
+
+        step = event.step
+        fail_count = 0
+        _err_info = None
+
+        for value in event.summary.value:
+            tag = value.tag
+            plugin_name = value.metadata.plugin_data.plugin_name
+            if len(plugin_name) > 0 and plugin_name not in self.supported_plugins:
+                if not self.unsupported_plugin_noticed:
+                    logging.warning(
+                        "Found unsupported plugin type({}) in the log file. "
+                        "Data for these wont be processed. "
+                        "Supported plugin types are: {}".format(plugin_name, ", ".join(self.supported_plugins)),
+                    )
+                    self.unsupported_plugin_noticed = True
+                continue
+            track_val = None
+            try:
+                if value.HasField("tensor"):
+                    # TODO: [MV] check the case when audios are passed via tensor
+                    if plugin_name == "images":
+                        tensor = value.tensor.string_val[2:]
+                        track_val = [Image(tf.image.decode_image(t).numpy()) for t in tensor]
+                        if len(track_val) == 1:
+                            track_val = track_val[0]
+                    elif plugin_name == "histograms":
+                        track_val = _decode_histogram_from_plugin(value)
+                    elif plugin_name == "scalars" or plugin_name == "":
+                        track_val = create_ndarray(value.tensor)
+                    else:
+                        track_val = value.tensor.float_val[0]
+                elif value.HasField("simple_value"):
+                    track_val = value.simple_value
+                elif value.HasField("image"):
+                    track_val = Image(tf.image.decode_image(value.image.encoded_image_string).numpy())
+                elif value.HasField("audio"):
+                    tf_audio, sample_rate = tf.audio.decode_wav(value.audio.encoded_audio_string)
+                    track_val = Audio(tf_audio.numpy(), rate=sample_rate)
+                elif value.HasField("histo"):
+                    track_val = _decode_histogram(value)
+
+            except RuntimeError as exc:
+                # catch all the nasty failures
+                fail_count += 1
+                if not _err_info:
+                    _err_info = str(exc)
+                continue
+
+            if track_val is not None:
+                self.queue.put(TensorboardEvent(track_val, tag, step, context={'entry': self.folder_name}))
+        if fail_count:
+            logging.warning(f"Failed to process {fail_count} entries. First exception: {_err_info}")
+
+
+class TensorboardEvent:
+
+    def __init__(self, value: Any, name: str, step: int, context: dict) -> None:
+        self.value = value
+        self.name = name
+        self.step = step
+        self.context = context
+
+
+class TensorboardEventConsumer:
+
+    def __init__(self, queue: queue.Queue, tracker) -> None:
+        self._tracker = weakref.ref(tracker)
+        self._queue = queue
+        self._thread = threading.Thread(target=self._process_events, daemon=True)
+        self._shutdown = False
+        self._started = False
+
+    def start(self):
+        if self._started:
+            return
+        self._started = True
+        self._thread.start()
+
+    def _process_events(self):
+        while True:
+            try:
+                event = self._queue.get(True, 1)
+                if event:
+                    self._tracker()(event.value, event.name, event.step, context=event.context)
+            except queue.Empty:
+                event = None
+                if self._shutdown:
+                    break
+
+    def stop(self):
+        if not self._started:
+            return
+        self._shutdown = True
+        self._thread.join()

aim/ext/transport/__init__.py

@@ -0,0 +1 @@
+from .client import Client

aim/ext/transport/client.py

@@ -0,0 +1,299 @@
+import requests
+import base64
+import logging
+import os
+import threading
+import uuid
+import weakref
+
+from copy import deepcopy
+from typing import Tuple
+from websockets.sync.client import connect
+
+from aim.ext.transport.utils import handle_exception
+from aim.ext.transport.message_utils import (
+    raise_exception,
+    pack_args,
+    unpack_stream,
+    unpack_args,
+    encode_tree,
+    decode_tree
+)
+from aim.ext.transport.request_queue import RequestQueue
+from aim.ext.transport.heartbeat import HeartbeatSender
+
+AIM_CLIENT_QUEUE_MAX_MEMORY = '__AIM_CLIENT_QUEUE_MAX_MEMORY__'
+DEFAULT_RETRY_INTERVAL = 0.1  # 100 ms
+DEFAULT_RETRY_COUNT = 2
+
+logger = logging.getLogger(__name__)
+
+
+class Client:
+    _thread_local = threading.local()
+
+    _queue = RequestQueue(
+        'remote_tracker',
+        max_queue_memory=os.getenv(AIM_CLIENT_QUEUE_MAX_MEMORY, 1024 * 1024 * 1024),
+        retry_count=DEFAULT_RETRY_COUNT,
+        retry_interval=DEFAULT_RETRY_INTERVAL
+    )
+
+    def __init__(self, remote_path: str):
+        # temporary workaround for M1 build
+
+        self._id = str(uuid.uuid4())
+        if remote_path.endswith('/'):
+            remote_path = remote_path[:-1]
+        self._remote_path = remote_path
+
+        self._http_protocol = 'http://'
+        self._ws_protocol = 'ws://'
+        self.protocol_probe()
+
+        self._resource_pool = weakref.WeakValueDictionary()
+
+        self._client_endpoint = f'{self.remote_path}/client'
+        self._tracking_endpoint = f'{self.remote_path}/tracking'
+        self.request_headers = {}
+        self.connect()
+
+        self._heartbeat_sender = HeartbeatSender(self)
+        self._heartbeat_sender.start()
+        self._thread_local.atomic_instructions = {}
+        self._ws = None
+
+    def protocol_probe(self):
+        endpoint = f'http://{self.remote_path}/status/'
+        try:
+            response = requests.get(endpoint, headers=self.request_headers)
+            if response.status_code == 200:
+                if response.url.startswith('https://'):
+                    self._http_protocol = 'https://'
+                    self._ws_protocol = 'wss://'
+                return
+        except Exception:
+            pass
+
+        endpoint = f'https://{self.remote_path}/status/'
+        try:
+            response = requests.get(endpoint, headers=self.request_headers)
+            if response.status_code == 200:
+                self._http_protocol = 'https://'
+                self._ws_protocol = 'wss://'
+        except Exception:
+            pass
+
+    def reinitialize_resource(self, handler):
+        # write some request to get a resource on server side with an already given handler
+        resource = self._resource_pool[handler]
+        self.get_resource_handler(resource, resource.resource_type, handler, resource.init_args)
+
+    def _reinitialize_all_resources(self):
+        handlers_list = list(self._resource_pool.keys())
+        for handler in handlers_list:
+            self.reinitialize_resource(handler)
+
+    @handle_exception(requests.ConnectionError,
+                      error_message='Failed to connect to Aim Server. Have you forgot to run `aim server` command?')
+    def _check_remote_version_compatibility(self):
+        from aim.__version__ import __version__ as client_version
+
+        error_message_template = 'The Aim Remote tracking server version ({}) '\
+                                 'is not compatible with the Aim client version ({}).'\
+                                 'Please upgrade either the Aim Client or the Aim Remote.'
+
+        warning_message_template = 'The Aim Remote tracking server version ({}) ' \
+                                   'and the Aim client version ({}) do not match.' \
+                                   'Consider upgrading either the client or remote tracking server.'
+
+        remote_version = self.get_version()
+
+        # server doesn't yet have the `get_version()` method implemented
+        if remote_version == '<3.19.0':
+            RuntimeError(error_message_template.format(remote_version, client_version))
+
+        # compare versions
+        if client_version == remote_version:
+            return
+
+        # if the server has a newer version always force to upgrade the client
+        if client_version < remote_version:
+            raise RuntimeError(error_message_template.format(remote_version, client_version))
+
+        # for other mismatching versions throw a warning for now
+        logger.warning(warning_message_template.format(remote_version, client_version))
+        # further incompatibility list will be added manually
+
+    def client_heartbeat(self):
+        endpoint = f'{self._http_protocol}{self._client_endpoint}/heartbeat/{self.uri}/'
+        response = requests.get(endpoint, headers=self.request_headers)
+        response_json = response.json()
+        if response.status_code != 200:
+            raise_exception(response_json.get('message'))
+
+        return response
+
+    @handle_exception(requests.ConnectionError,
+                      error_message='Failed to connect to Aim Server. Have you forgot to run `aim server` command?')
+    def connect(self):
+        endpoint = f'{self._http_protocol}{self._client_endpoint}/connect/{self.uri}/'
+        response = requests.get(endpoint, headers=self.request_headers)
+        response_json = response.json()
+        if response.status_code != 200:
+            raise_exception(response_json.get('message'))
+
+        return response
+
+    def reconnect(self):
+        endpoint = f'{self._http_protocol}{self._client_endpoint}/reconnect/{self.uri}/'
+        response = requests.get(endpoint, headers=self.request_headers)
+        response_json = response.json()
+        if response.status_code != 200:
+            raise_exception(response_json.get('message'))
+
+        self.refresh_ws()
+        self._reinitialize_all_resources()
+
+        return response
+
+    def disconnect(self):
+        self._heartbeat_sender.stop()
+        if self._ws:
+            self._ws.close()
+
+        endpoint = f'{self._http_protocol}{self._client_endpoint}/disconnect/{self.uri}/'
+        response = requests.get(endpoint, headers=self.request_headers)
+        response_json = response.json()
+        if response.status_code != 200:
+            raise_exception(response_json.get('message'))
+
+        return response
+
+    def get_version(self,):
+        endpoint = f'{self._http_protocol}{self._client_endpoint}/get-version/'
+        response = requests.get(endpoint, headers=self.request_headers)
+        response_json = response.json()
+        if response.status_code == 404:
+            return '<3.19.0'
+        if response.status_code == 400:
+            raise_exception(response_json.get('exception'))
+        return response_json.get('version')
+
+    def get_resource_handler(self, resource, resource_type, handler='', args=()):
+        endpoint = f'{self._http_protocol}{self._tracking_endpoint}/{self.uri}/get-resource/'
+
+        request_data = {
+            'resource_handler': handler,
+            'resource_type': resource_type,
+            'args': base64.b64encode(args).decode()
+        }
+
+        response = requests.post(endpoint, json=request_data, headers=self.request_headers)
+        response_json = response.json()
+        if response.status_code == 400:
+            raise_exception(response_json.get('exception'))
+        elif response.status_code != 200:
+            raise (Exception(response_json))
+
+        handler = response_json.get('handler')
+        self._resource_pool[handler] = resource
+
+        return handler
+
+    def release_resource(self, queue_id, resource_handler):
+        endpoint = f'{self._http_protocol}{self._tracking_endpoint}/{self.uri}/release-resource/{resource_handler}/'
+        if queue_id != -1:
+            self.get_queue().wait_for_finish()
+
+        response = requests.get(endpoint, headers=self.request_headers)
+        response_json = response.json()
+        if response.status_code == 400:
+            raise_exception(response_json.get('exception'))
+
+        del self._resource_pool[resource_handler]
+
+    def run_instruction(self, queue_id, resource, method, args=(), is_write_only=False):
+        args = deepcopy(args)
+
+        # self._thread_local can be empty in the 'clean up' phase.
+
+        if is_write_only:
+            assert queue_id != -1
+            if getattr(self._thread_local, 'atomic_instructions', None) is not None and \
+                    self._thread_local.atomic_instructions.get(queue_id, None) is not None:
+                self._thread_local.atomic_instructions[queue_id].append((resource, method, args))
+                return
+
+            self.get_queue().register_task(
+                self,
+                self._run_write_instructions, list(encode_tree([(resource, method, args)], strict=False)))
+            return
+
+        return self._run_read_instructions(queue_id, resource, method, args)
+
+    def _run_read_instructions(self, queue_id, resource, method, args):
+        endpoint = f'{self._http_protocol}{self._tracking_endpoint}/{self.uri}/read-instruction/'
+
+        request_data = {
+            'resource_handler': resource,
+            'method_name': method,
+            'args': base64.b64encode(pack_args(encode_tree(args))).decode()
+        }
+
+        if queue_id != -1:
+            self.get_queue().wait_for_finish()
+
+        response = requests.post(endpoint, json=request_data, stream=True, headers=self.request_headers)
+
+        if response.status_code == 400:
+            raise_exception(response.json().get('exception'))
+        return decode_tree(unpack_stream(response.iter_content(chunk_size=None)))
+
+    def _run_write_instructions(self, instructions: [Tuple[bytes, bytes]]):
+        msg = pack_args(iter(instructions))
+
+        self.ws.send(msg)
+
+        response = self.ws.recv()
+        if response == b'OK':
+            return
+
+        response_json = decode_tree(unpack_args(response))
+        raise_exception(response_json)
+
+    def start_instructions_batch(self, hash_):
+        if getattr(self._thread_local, 'atomic_instructions', None) is None:
+            self._thread_local.atomic_instructions = {}
+        self._thread_local.atomic_instructions[hash_] = []
+
+    def flush_instructions_batch(self, hash_):
+        if self._thread_local.atomic_instructions.get(hash_) is None:
+            return
+
+        self.get_queue().register_task(
+            self,
+            self._run_write_instructions, list(encode_tree(self._thread_local.atomic_instructions[hash_])))
+        del self._thread_local.atomic_instructions[hash_]
+
+    def refresh_ws(self):
+        self._ws = connect(f'{self._ws_protocol}{self._tracking_endpoint}/{self.uri}/write-instruction/')
+
+    @property
+    def ws(self):
+        if self._ws is None:
+            self._ws = connect(f'{self._ws_protocol}{self._tracking_endpoint}/{self.uri}/write-instruction/',
+                               additional_headers=self.request_headers)
+
+        return self._ws
+
+    @property
+    def uri(self):
+        return self._id
+
+    @property
+    def remote_path(self):
+        return self._remote_path
+
+    def get_queue(self):
+        return self._queue

aim/ext/transport/config.py

@@ -0,0 +1,7 @@
+# AIM SERVER
+AIM_SERVER_DEFAULT_PORT = 53800
+AIM_SERVER_DEFAULT_HOST = '0.0.0.0'
+AIM_SERVER_MOUNTED_REPO_PATH = '__AIM_SERVER_MOUNT_REPO_PATH__'
+AIM_SERVER_BASE_PATH = '__AIM_SERVER_BASE_PATH__'
+
+AIM_RT_BEARER_TOKEN = '__AIM_RT_BEARER_TOKEN__'

aim/ext/transport/handlers.py

@@ -0,0 +1,114 @@
+import os
+import uuid
+import pathlib
+import pytz
+
+from datetime import datetime
+
+from aim.ext.transport.config import AIM_SERVER_MOUNTED_REPO_PATH
+
+from aim.sdk import Repo
+from aim.sdk.reporter import RunStatusReporter, ScheduledStatusReporter
+from aim.sdk.reporter.file_manager import LocalFileManager
+from aim.ext.cleanup import AutoClean
+
+
+class ResourceRefAutoClean(AutoClean['ResourceRef']):
+    @staticmethod
+    def noop(res: object):
+        return
+
+    def __init__(self, instance: 'ResourceRef'):
+        super().__init__(instance)
+        self._finalizer_func = instance._finalizer_func
+        self._resource = instance._resource
+
+    def _close(self):
+        self._finalizer_func(self._resource)
+
+
+class ResourceRef:
+    def __init__(self, res_obj, finalizer_func=ResourceRefAutoClean.noop):
+        self._resource = res_obj
+        self._finalizer_func = finalizer_func
+        self._auto_clean = ResourceRefAutoClean(self)
+
+    @property
+    def ref(self):
+        return self._resource
+
+
+def get_handler():
+    return str(uuid.uuid4())
+
+
+def get_tree(**kwargs):
+    repo_path = os.environ.get(AIM_SERVER_MOUNTED_REPO_PATH)
+    if repo_path:
+        repo = Repo.from_path(repo_path)
+    else:
+        repo = Repo.default_repo()
+    name = kwargs['name']
+    sub = kwargs['sub']
+    read_only = kwargs['read_only']
+    from_union = kwargs['from_union']
+    index = kwargs['index']
+    timeout = kwargs['timeout']
+    no_cache = kwargs.get('no_cache', False)
+    if index:
+        return ResourceRef(repo._get_index_tree(name, timeout))
+    else:
+        return ResourceRef(repo.request_tree(name, sub, read_only=read_only, from_union=from_union, no_cache=no_cache))
+
+
+def get_structured_run(hash_, read_only, created_at, **kwargs):
+    repo_path = os.environ.get(AIM_SERVER_MOUNTED_REPO_PATH)
+    if repo_path:
+        repo = Repo.from_path(repo_path)
+    else:
+        repo = Repo.default_repo()
+    if created_at is not None:
+        created_at = datetime.fromtimestamp(created_at, tz=pytz.utc).replace(tzinfo=None)
+    return ResourceRef(repo.request_props(hash_, read_only, created_at))
+
+
+def get_repo():
+    repo_path = os.environ.get(AIM_SERVER_MOUNTED_REPO_PATH)
+    if repo_path:
+        repo = Repo.from_path(repo_path)
+    else:
+        repo = Repo.default_repo()
+    return ResourceRef(repo)
+
+
+def get_lock(**kwargs):
+    repo_path = os.environ.get(AIM_SERVER_MOUNTED_REPO_PATH)
+    if repo_path:
+        repo = Repo.from_path(repo_path)
+    else:
+        repo = Repo.default_repo()
+    run_hash = kwargs['run_hash']
+    # TODO Do we need to import SFRunLock here?
+    from aim.sdk.lock_manager import SFRunLock
+    return ResourceRef(repo.request_run_lock(run_hash), SFRunLock.release)
+
+
+def get_run_heartbeat(run_hash, **kwargs):
+    repo_path = os.environ.get(AIM_SERVER_MOUNTED_REPO_PATH)
+    if repo_path:
+        repo = Repo.from_path(repo_path)
+    else:
+        repo = Repo.default_repo()
+    status_reporter = RunStatusReporter(run_hash, LocalFileManager(repo.path))
+    progress_flag_path = pathlib.Path(repo.path) / 'meta' / 'progress' / run_hash
+    return ResourceRef(ScheduledStatusReporter(status_reporter, touch_path=progress_flag_path),
+                       ScheduledStatusReporter.stop)
+
+
+def get_file_manager(**kwargs):
+    repo_path = os.environ.get(AIM_SERVER_MOUNTED_REPO_PATH)
+    if repo_path:
+        repo = Repo.from_path(repo_path)
+    else:
+        repo = Repo.default_repo()
+    return ResourceRef(LocalFileManager(repo.path))

aim/ext/transport/heartbeat.py

@@ -0,0 +1,163 @@
+import datetime
+import logging
+import time
+import weakref
+
+from collections import Counter
+from threading import Thread
+from typing import Union
+
+logger = logging.getLogger(__name__)
+
+
+class HeartbeatSender(object):
+
+    HEARTBEAT_INTERVAL_DEFAULT = 10
+    NETWORK_CHECK_INTERVAL = 180
+
+    NETWORK_UNSTABLE_WARNING_TEMPLATE = 'Network connection between client `{}` ' \
+                                        'and server `{}` appears to be unstable.'
+    NETWORK_ABSENT_WARNING_TEMPLATE = 'Network connection between client `{}` ' \
+                                      'and server `{}` appears to be absent.'
+
+    def __init__(self,
+                 client,
+                 interval: Union[int, float] = HEARTBEAT_INTERVAL_DEFAULT,
+                 ):
+        self._remote_client = weakref.ref(client)
+        self._heartbeat_send_interval = interval
+
+        # network state check vars
+        self._network_stability_check_interval = HeartbeatSender.NETWORK_CHECK_INTERVAL
+        self._network_unstable_warned = False
+        self._network_absent_warned = False
+        self._heartbeat_responses = Counter(success=0, fail=0)
+
+        # Start thread to collect stats and logs at intervals
+        self._th_collector = Thread(target=self._target_f, daemon=True)
+        self._shutdown = False
+        self._started = False
+
+    def start(self):
+        if self._started:
+            return
+
+        self._started = True
+        self._th_collector.start()
+
+    def stop(self):
+        if not self._started:
+            return
+
+        self._shutdown = True
+        self._th_collector.join()
+
+    def _send_heartbeat(self):
+        if self._remote_client():
+            try:
+                response = self._remote_client().client_heartbeat()
+                if response.status_code == 200:
+                    self._heartbeat_responses['success'] += 1
+                else:
+                    self._heartbeat_responses['fail'] += 1
+            except Exception:
+                # at the moment we don't care about failures for heartbeats
+                self._heartbeat_responses['fail'] += 1
+
+    def _target_f(self):
+        heartbeat_interval_counter = 0
+        stability_check_interval_counter = 0
+        # send initial heartbeat
+        self._send_heartbeat()
+
+        while True:
+            if self._shutdown:
+                break
+
+            time.sleep(1)
+            heartbeat_interval_counter += 1
+            stability_check_interval_counter += 1
+
+            if heartbeat_interval_counter > self._heartbeat_send_interval:
+                self._send_heartbeat()
+                heartbeat_interval_counter = 0
+
+            if stability_check_interval_counter > self._network_stability_check_interval:
+                self._check_network_state()
+                stability_check_interval_counter = 0
+
+    def _check_network_state(self):
+        def reset_responses():
+            self._heartbeat_responses['fail'] = 0
+            self._heartbeat_responses['success'] = 0
+
+        if not self._heartbeat_responses['fail']:
+            reset_responses()
+            return
+
+        if self._heartbeat_responses['success'] and not self._network_unstable_warned:
+            self._network_unstable_warned = True
+            logger.warning(HeartbeatSender.NETWORK_UNSTABLE_WARNING_TEMPLATE
+                           .format(self._remote_client().uri, self._remote_client().remote_path))
+            reset_responses()
+            return
+
+        if not self._network_absent_warned:
+            self._network_absent_warned = True
+            logger.warning(HeartbeatSender.NETWORK_ABSENT_WARNING_TEMPLATE
+                           .format(self._remote_client().uri, self._remote_client().remote_path))
+
+        reset_responses()
+
+
+class HeartbeatWatcher:
+    CLIENT_KEEP_ALIVE_TIME_DEFAULT = 30 * 60  # 30 minutes
+
+    def __init__(self,
+                 heartbeat_pool,
+                 keep_alive_time: Union[int, float] = CLIENT_KEEP_ALIVE_TIME_DEFAULT):
+
+        self._heartbeat_pool = heartbeat_pool
+
+        self._client_keep_alive_time = keep_alive_time
+
+        # Start thread to collect stats and logs at intervals
+        self._th_collector = Thread(target=self._interval_check, daemon=True)
+        self._shutdown = False
+        self._started = False
+
+    def start(self):
+        if self._started:
+            return
+
+        self._started = True
+        self._th_collector.start()
+
+    def stop(self):
+        if not self._started:
+            return
+
+        self._shutdown = True
+        self._th_collector.join()
+
+    def _release_client_resources(self, dead_client_uri):
+        logger.warning(f'Cleaning up resources for client `{dead_client_uri}`.')
+        from aim.ext.transport.tracking import TrackingRouter
+        from aim.ext.transport.router import ClientRouter
+        TrackingRouter.cleanup_client_resources(dead_client_uri)
+        ClientRouter.remove_client(dead_client_uri)
+
+    def _interval_check(self):
+        while True:
+            # Get system statistics
+            if self._shutdown:
+                break
+
+            time.sleep(1)
+
+            client_uris = list(self._heartbeat_pool.keys())
+            for client_uri in client_uris:
+                now = datetime.datetime.now().timestamp()
+                if now - self._heartbeat_pool[client_uri] > self._client_keep_alive_time:
+                    self._release_client_resources(client_uri)
+                    del self._heartbeat_pool[client_uri]

aim/ext/transport/message_utils.py

@@ -0,0 +1,85 @@
+import importlib
+import struct
+import json
+from typing import Iterator, Tuple
+
+from aim.storage.object import CustomObject
+from aim.storage.types import BLOB
+from aim.storage.treeutils import encode_tree, decode_tree # noqa
+
+
+def pack_args(tree: Iterator[Tuple[bytes, bytes]]) -> bytes:
+    result = []
+    for key, val in tree:
+        if not isinstance(val, BLOB):
+            result.append(struct.pack('I', len(key)) + key + struct.pack('?', False) + struct.pack('I', len(val)) + val)
+        else:
+            val = val.load()
+            result.append(struct.pack('I', len(key)) + key + struct.pack('?', True) + struct.pack('I', len(val)) + val)
+
+    return b''.join(result)
+
+
+def unpack_args(args: bytes) -> Tuple[bytes, bytes]:
+    while args:
+        (key_size,), args_tail = struct.unpack('I', args[:4]), args[4:]
+        key, args_tail = args_tail[:key_size], args_tail[key_size:]
+        (is_blob,), args_tail = struct.unpack('?', args_tail[:1]), args_tail[1:]
+        (value_size,), args_tail = struct.unpack('I', args_tail[:4]), args_tail[4:]
+        value, args_tail = args_tail[:value_size], args_tail[value_size:]
+        args = args_tail
+        if is_blob:
+            yield key, BLOB(data=value)
+        else:
+            yield key, value
+
+
+def pack_stream(tree: Iterator[Tuple[bytes, bytes]]) -> bytes:
+    # TODO: [MV] check the performance diff of current version vs collecting the whole tree as a chunk
+    for key, val in tree:
+        if not isinstance(val, BLOB):
+            yield struct.pack('I', len(key)) + key + struct.pack('?', False) + struct.pack('I', len(val)) + val
+        else:
+            val = val.load()
+            yield struct.pack('I', len(key)) + key + struct.pack('?', True) + struct.pack('I', len(val)) + val
+
+
+def unpack_helper(msg: bytes) -> Tuple[bytes, bytes]:
+    (key_size,), tail = struct.unpack('I', msg[:4]), msg[4:]
+    key, tail = tail[:key_size], tail[key_size:]
+    (is_blob,), tail = struct.unpack('?', tail[:1]), tail[1:]
+    (value_size,), tail = struct.unpack('I', tail[:4]), tail[4:]
+    value, tail = tail[:value_size], tail[value_size:]
+    assert len(tail) == 0
+    if is_blob:
+        yield key, BLOB(data=value)
+    else:
+        yield key, value
+
+
+def unpack_stream(stream) -> Tuple[bytes, bytes]:
+    for msg in stream:
+        yield from unpack_helper(msg)
+
+
+def raise_exception(server_exception):
+    module = importlib.import_module(server_exception.get('module_name'))
+    exception = getattr(module, server_exception.get('class_name'))
+    args = json.loads(server_exception.get('args') or [])
+    raise exception(*args) if args else exception()
+
+
+def build_exception(exception: Exception):
+    return {
+        'module_name': exception.__class__.__module__,
+        'class_name': exception.__class__.__name__,
+        'args': json.dumps(exception.args),
+    }
+
+
+@CustomObject.alias('aim.resource')
+class ResourceObject(CustomObject):
+    AIM_NAME = 'aim.resource'
+
+    def __init__(self, handler):
+        self.storage['handler'] = handler

aim/ext/transport/remote_resource.py

@@ -0,0 +1,15 @@
+from aim.ext.cleanup import AutoClean
+from .client import Client
+
+
+class RemoteResourceAutoClean(AutoClean):
+    def __init__(self, instance):
+        super().__init__(instance)
+        self.hash = -1
+        self.handler = None
+        self.rpc_client: Client = None
+
+    def _close(self):
+        if self.handler is not None:
+            assert self.rpc_client is not None
+            self.rpc_client.release_resource(self.hash, self.handler)

aim/ext/transport/request_queue.py

@@ -0,0 +1,117 @@
+import time
+import queue
+import logging
+import threading
+import weakref
+
+logger = logging.getLogger(__name__)
+
+
+class RequestQueue(object):
+    def __init__(self, name, max_queue_memory=0,
+                 retry_count=0, retry_interval=0):
+
+        self._client = None
+
+        self.retry_count = retry_count or 1
+        self.retry_interval = retry_interval
+        self._needs_reconnect = False
+
+        self.max_memory_usage = max_queue_memory
+        self.current_memory_usage = 0
+
+        self._shutdown = False
+        self._queue = queue.Queue()
+        self._name = name
+
+        self._thread = threading.Thread(target=self.worker)
+        self._thread.daemon = True
+        self._thread.start()
+
+    def register_task(self, client, task_f, *args):
+        if not self._client:
+            self._client = weakref.ref(client)
+
+        if self._shutdown:
+            logger.debug('Cannot register task: rpc task queue is stopped.')
+            return
+
+        arg_size = self._calculate_size(args)
+        with self._queue.not_full:
+            while self.current_memory_usage + arg_size >= self.max_memory_usage:
+                self._queue.not_full.wait()
+
+        with self._queue.mutex:
+            self.current_memory_usage += arg_size
+
+        self._queue.put((task_f, args))
+
+    def worker(self):
+        while True:
+            if self._shutdown:
+                logger.debug(f'Shutting down worker thread {threading.get_ident()}.')
+                break
+            task_f, args = self._queue.get()
+            if self._try_exec_task(task_f, *args):
+                arg_size = self._calculate_size(args)
+                with self._queue.mutex:
+                    self.current_memory_usage -= arg_size
+                # clear the unnecessary references
+                task_f, args = None, None
+                self._queue.task_done()
+            else:
+                self._put_front(task_f, args)
+
+    def _try_exec_task(self, task_f, *args):
+        # temporary workaround for M1 build
+        from websockets.exceptions import ConnectionClosedError
+
+        retry = 0
+        while retry < self.retry_count:
+            if self._needs_reconnect:
+                try:
+                    self._client().reconnect()
+                    self._needs_reconnect = False
+                except Exception:
+                    retry += 1
+                    time.sleep(self.retry_interval)
+                    continue
+
+            try:
+                task_f(*args)
+                return True
+            except ConnectionClosedError as e:
+                self._needs_reconnect = True
+
+                retry += 1
+                time.sleep(self.retry_interval)
+                logger.warning(f'Remote Server is unavailable, please check network connection: {e}.')
+
+        return False
+
+    def _put_front(self, task_f, args):
+        with self._queue.not_full:
+            self._queue.queue.appendleft((task_f, args))
+            self._queue.not_empty.notify()
+
+    def wait_for_finish(self):
+        self._queue.join()
+
+    def stop(self):
+        pending_task_count = self._queue.qsize()
+        if pending_task_count:
+            logger.warning(f'Processing {pending_task_count} pending tasks in the rpc queue \'{self._name}\'... '
+                           f'Please do not kill the process.')
+            self._queue.join()
+        logger.debug('No pending tasks left.')
+        self._shutdown = True
+
+    @staticmethod
+    def _calculate_size(args):
+        size = 0
+        assert type(args) is tuple
+        for arg in args[0]:
+            assert type(arg) is tuple
+            assert len(arg) == 2
+            size += len(arg[0]) + len(arg[1])
+        return size

aim/ext/transport/router.py

@@ -0,0 +1,43 @@
+import datetime
+
+from fastapi import APIRouter
+
+
+class ClientRouter:
+    client_heartbeat_pool = dict()
+    clients = []
+
+    def __init__(self):
+        self.router = APIRouter()
+        self.router.add_api_route('/get-version/', self.get_version, methods=['GET'])
+        self.router.add_api_route('/heartbeat/{client_uri}/', self.heartbeat, methods=['GET'])
+        self.router.add_api_route('/connect/{client_uri}/', self.connect, methods=['GET'])
+        self.router.add_api_route('/reconnect/{client_uri}/', self.reconnect, methods=['GET'])
+        self.router.add_api_route('/disconnect/{client_uri}/', self.disconnect, methods=['GET'])
+
+    @classmethod
+    def add_client(cls, client_uri):
+        cls.clients.append(client_uri)
+
+    @classmethod
+    def remove_client(cls, client_uri):
+        if client_uri in cls.clients:
+            cls.clients.remove(client_uri)
+
+    async def get_version(self):
+        from aim.__version__ import __version__ as aim_version
+
+        return {'version': aim_version}
+
+    async def heartbeat(self, client_uri):
+        self.client_heartbeat_pool[client_uri] = datetime.datetime.now().timestamp()
+
+    async def connect(self, client_uri):
+        self.add_client(client_uri)
+
+    async def reconnect(self, client_uri):
+        if client_uri not in self.clients:
+            self.add_client(client_uri)
+
+    async def disconnect(self, client_uri):
+        self.remove_client(client_uri)

aim/ext/transport/run.py

@@ -0,0 +1,3 @@
+from aim.ext.transport.server import create_app
+
+app = create_app()

aim/ext/transport/server.py

@@ -0,0 +1,71 @@
+import os
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.exceptions import HTTPException
+
+from aim.ext.utils import http_exception_handler, fallback_exception_handler
+from aim.ext.transport.router import ClientRouter
+from aim.ext.transport.tracking import TrackingRouter, ResourceTypeRegistry
+from aim.ext.transport.heartbeat import HeartbeatWatcher
+from aim.ext.transport.handlers import (
+    get_tree,
+    get_structured_run,
+    get_repo,
+    get_lock,
+    get_run_heartbeat,
+    get_file_manager
+)
+from aim.ext.transport.config import AIM_SERVER_BASE_PATH
+
+
+def prepare_resource_registry():
+    registry = ResourceTypeRegistry()
+    registry.register('TreeView', get_tree)
+    registry.register('StructuredRun', get_structured_run)
+    registry.register('Repo', get_repo)
+    registry.register('Lock', get_lock)
+    registry.register('RunHeartbeat', get_run_heartbeat)
+    registry.register('FileManager', get_file_manager)
+    return registry
+
+
+def create_app():
+    app = FastAPI(title=__name__)
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=['*'],
+        allow_methods=['GET', 'POST', 'PUT', 'DELETE', 'OPTIONS', 'HEAD'],
+        allow_headers=['Origin', 'X-Requested-With',
+                       'Content-Type', 'Accept', 'Authorization', 'X-Timezone-Offset'],
+        allow_credentials=True,
+        max_age=86400
+    )
+
+    registry = prepare_resource_registry()
+
+    client_router = ClientRouter()
+    tracking_router = TrackingRouter(registry)
+
+    watcher = HeartbeatWatcher(
+        ClientRouter.client_heartbeat_pool,
+    )
+    watcher.start()
+
+    api_app = FastAPI()
+
+    api_app.add_exception_handler(HTTPException, http_exception_handler)
+    api_app.add_exception_handler(Exception, fallback_exception_handler)
+
+    api_app.include_router(client_router.router, prefix='/client')
+    api_app.include_router(tracking_router.router, prefix='/tracking')
+
+    base_path = os.environ.get(AIM_SERVER_BASE_PATH, '')
+    app.mount(f'{base_path}/', api_app)
+
+    @api_app.get('/status/')
+    async def status():
+        return {'status': 'OK'}
+
+    return app

aim/ext/transport/tracking.py

@@ -0,0 +1,221 @@
+import uuid
+import base64
+import logging
+
+from typing import Dict, Union, List
+from fastapi import WebSocket, Request, APIRouter, WebSocketDisconnect
+from fastapi.responses import StreamingResponse, JSONResponse
+
+from aim.ext.transport.message_utils import (
+    encode_tree,
+    decode_tree,
+    unpack_args,
+    pack_args,
+    pack_stream,
+    build_exception,
+    ResourceObject,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def get_handler():
+    return str(uuid.uuid4())
+
+
+class UnauthorizedRequestError(RuntimeError):
+    def __init__(self, handler, *args, **kwargs):
+        self.handler = handler
+
+
+class ResourceTypeRegistry:
+    def __init__(self):
+        self._registry: Dict[str, type] = {}
+
+    def register(self, type_name: str, resource_getter: Union[type, callable]):
+        self._registry[type_name] = resource_getter
+
+    def __getitem__(self, type_name: str):
+        return self._registry[type_name]
+
+
+class ConnectionManager:
+    def __init__(self):
+        self.active_connections: List[WebSocket] = []
+
+    async def connect(self, websocket: WebSocket):
+        await websocket.accept()
+        self.active_connections.append(websocket)
+
+    def disconnect(self, websocket: WebSocket):
+        self.active_connections.remove(websocket)
+
+    async def send_personal_message(self, message: str, websocket: WebSocket):
+        await websocket.send_text(message)
+
+    async def broadcast(self, message: str):
+        for connection in self.active_connections:
+            await connection.send_text(message)
+
+
+class TrackingRouter:
+    resource_pool = dict()
+    manager = ConnectionManager()
+
+    def __init__(self, resource_registry: ResourceTypeRegistry):
+        self.registry = resource_registry
+        self.router = APIRouter()
+        self.router.add_api_route('/{client_uri}/get-resource/',
+                                  self.get_resource, methods=['POST'])
+        self.router.add_api_route('/{client_uri}/release-resource/{resource_handler}/',
+                                  self.release_resource, methods=['GET'])
+        self.router.add_api_route('/{client_uri}/read-instruction/', self.run_instruction, methods=['POST'])
+        self.router.add_api_websocket_route('/{client_uri}/write-instruction/', self.run_write_instructions)
+
+    @classmethod
+    def cleanup_client_resources(cls, dead_client_uri):
+        resource_handlers = list(cls.resource_pool.keys())
+        for handler in resource_handlers:
+            (client_uri, _) = cls.resource_pool[handler]
+            if dead_client_uri == client_uri:
+                del cls.resource_pool[handler]
+
+    @classmethod
+    def _verify_resource_handler(cls, resource_handler, client_uri):
+        res_info = cls.resource_pool.get(resource_handler, None)
+        if not res_info or res_info[0] != client_uri:
+            raise UnauthorizedRequestError(resource_handler)
+
+    async def get_resource(self,
+                           client_uri: str,
+                           request: Request,):
+        request_data = await request.json()
+        resource_handler = request_data.get('resource_handler')
+        resource_type = request_data.get('resource_type')
+        args = request_data.get('args')
+
+        if not resource_handler:
+            resource_handler = get_handler()
+
+        try:
+            resource_cls = self.registry[resource_type]
+            if len(args) > 0:
+                kwargs = decode_tree(unpack_args(base64.b64decode(args)))
+                checked_kwargs = {}
+                for argname, arg in kwargs.items():
+                    if isinstance(arg, ResourceObject):
+                        handler = arg.storage['handler']
+                        self._verify_resource_handler(handler, client_uri)
+                        checked_kwargs[argname] = self.resource_pool[handler][1].ref
+                    else:
+                        checked_kwargs[argname] = arg
+
+                res = resource_cls(**checked_kwargs)
+            else:
+                res = resource_cls()
+
+            self.resource_pool[resource_handler] = (client_uri, res)
+            return {'handler': resource_handler}
+
+        except Exception as e:
+            try:
+                # clean up resource_pool
+                # because no one will call release_resource in case of Exception
+                del self.resource_pool[resource_handler]
+            except KeyError:
+                pass
+
+            logger.debug(f'Caught exception {e}. Sending response 400.')
+            return JSONResponse({
+                'exception': build_exception(e),
+            }, status_code=400)
+
+    async def release_resource(self, client_uri, resource_handler):
+        try:
+            self._verify_resource_handler(resource_handler, client_uri)
+            del self.resource_pool[resource_handler]
+        except Exception as e:
+            logger.debug(f'Caught exception {e}. Sending response 400.')
+            return JSONResponse({
+                'exception': build_exception(e),
+            }, status_code=400)
+
+    async def run_instruction(self, client_uri: str,
+                              request: Request,):
+        try:
+            request_data = await request.json()
+            resource_handler = request_data.get('resource_handler')
+            method_name = request_data.get('method_name')
+            args = request_data.get('args')
+
+            self._verify_resource_handler(resource_handler, client_uri)
+
+            args = decode_tree(unpack_args(base64.b64decode(args)))
+
+            checked_args = []
+            for arg in args:
+                if isinstance(arg, ResourceObject):
+                    handler = arg.storage['handler']
+                    self._verify_resource_handler(handler, client_uri)
+                    checked_args.append(self.resource_pool[handler][1].ref)
+                else:
+                    checked_args.append(arg)
+
+            resource = self.resource_pool[resource_handler][1].ref
+            if method_name.endswith('.setter'):
+                attr_name = method_name.split('.')[0]
+                setattr(resource, attr_name, checked_args[0])
+                result = None
+            else:
+                attr = getattr(resource, method_name)
+                if callable(attr):
+                    result = attr(*checked_args)
+                else:
+                    result = attr
+            del resource
+            return StreamingResponse(pack_stream(encode_tree(result)))
+        except Exception as e:
+            logger.debug(f'Caught exception {e}. Sending response 400.')
+            return JSONResponse({
+                'exception': build_exception(e),
+            }, status_code=400)
+
+    async def run_write_instructions(self, websocket: WebSocket, client_uri: str):
+        await self.manager.connect(websocket)
+
+        try:
+            while True:
+                raw_message = await websocket.receive_bytes()
+                write_instructions = decode_tree(
+                    unpack_args(raw_message))
+                for instruction in write_instructions:
+                    resource_handler, method_name, args = instruction
+                    self._verify_resource_handler(
+                        resource_handler, client_uri)
+                    checked_args = []
+                    for arg in args:
+                        if isinstance(arg, ResourceObject):
+                            handler = arg.storage['handler']
+                            self._verify_resource_handler(
+                                handler, client_uri)
+                            checked_args.append(
+                                self.resource_pool[handler][1].ref)
+                        else:
+                            checked_args.append(arg)
+
+                    resource = self.resource_pool[resource_handler][1].ref
+                    if method_name.endswith('.setter'):
+                        attr_name = method_name.split('.')[0]
+                        setattr(resource, attr_name, checked_args[0])
+                    else:
+                        attr = getattr(resource, method_name)
+                        assert callable(attr)
+                        attr(*checked_args)
+                    del resource
+
+                await websocket.send_bytes(b'OK')
+        except WebSocketDisconnect:
+            self.manager.disconnect(websocket)
+
+        except Exception as e:
+            await websocket.send_bytes(pack_args(encode_tree(build_exception(e))))

aim/ext/transport/utils.py

@@ -0,0 +1,22 @@
+import logging
+
+from functools import wraps
+from typing import Type, Optional
+
+logger = logging.getLogger(__name__)
+
+
+def handle_exception(exc_type: Type[Exception], error_message: Optional[str] = None):
+    def inner(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except exc_type as e:  # noqa
+                if error_message is not None:
+                    logger.error(error_message)
+                    raise RuntimeError(error_message)
+                else:  # silent mode
+                    pass
+        return wrapper
+    return inner

aim/ext/utils.py

@@ -0,0 +1,98 @@
+import logging
+import subprocess
+from fastapi.responses import JSONResponse
+
+logger = logging.getLogger(__name__)
+
+
+def get_installed_packages():
+    import pkg_resources
+    packages = {i.key: i.version for i in pkg_resources.working_set}
+
+    return packages
+
+
+def get_environment_variables():
+    from os import environ
+
+    env_mask = ('secret', 'key', 'token', 'password')
+    env_vars = {
+        k: v for k, v in environ.items() if next(
+            (m for m in env_mask if m in k.lower()), None
+        ) is None
+    }
+
+    return env_vars
+
+
+def get_git_info():
+    git_info = {}
+    try:
+        r = subprocess.run(['git', 'rev-parse', '--is-inside-work-tree'],
+                           stdout=subprocess.PIPE,
+                           stderr=subprocess.STDOUT,
+                           check=True)
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        # not a git repo
+        return git_info
+    else:
+        output = r.stdout.decode('utf-8').strip().lower()
+        if output != 'true':
+            # malformed result
+            return git_info
+
+    cmds = {
+        'branch': ('git', 'rev-parse', '--abbrev-ref', 'HEAD'),
+        'remote_origin_url': ('git', 'config', '--get', 'remote.origin.url'),
+        'commit': ('git', 'log', '--pretty=format:%h/%ad/%an', '--date=iso-strict', '-1'),
+    }
+    results = {}
+    for key, cmd in cmds.items():
+        try:
+            r = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, check=True)
+        except subprocess.CalledProcessError:
+            continue
+        else:
+            output = r.stdout.decode('utf-8').strip()
+            results[key] = output
+
+    try:
+        commit_hash, commit_timestamp, commit_author = results.get('commit').split('/')
+    except (ValueError, AttributeError):
+        commit_hash = commit_timestamp = commit_author = None
+
+    git_info.update({
+        'branch': results.get('branch'),
+        'remote_origin_url': results.get('remote_origin_url'),
+        'commit': {
+            'hash': commit_hash,
+            'timestamp': commit_timestamp,
+            'author': commit_author
+        }
+    })
+
+    return git_info
+
+
+async def http_exception_handler(request, exc):
+    message = str(exc.detail)
+    detail = None
+
+    if isinstance(exc.detail, dict):
+        message = exc.detail.pop('message', message)
+        detail = exc.detail.pop('detail', None)
+
+    response = {'message': message}
+    if detail:
+        response.update({'detail': detail})
+    else:
+        response.update({'detail': str(exc)})
+    return JSONResponse(response, status_code=exc.status_code)
+
+
+async def fallback_exception_handler(request, exc):
+    response = {
+        'message': f'\'{type(exc)}\' exception raised!',
+        'detail': str(exc)
+    }
+    return JSONResponse(response, status_code=500)

aim/fastai.py

@@ -0,0 +1,2 @@
+# Alias to SDK fast.ai interface
+from aim.sdk.adapters.fastai import AimCallback # noqa F401

aim/hf_dataset.py

@@ -0,0 +1,2 @@
+# Alias to SDK Hugging Face Datasets interface
+from aim.sdk.objects.plugins.hf_datasets_metadata import HFDataset  # noqa F401

aim/hugging_face.py

@@ -0,0 +1,2 @@
+# Alias to SDK Hugging Face interface
+from aim.sdk.adapters.hugging_face import AimCallback  # noqa F401

aim/keras.py

@@ -0,0 +1,2 @@
+# Alias to SDK Keras interface
+from aim.sdk.adapters.keras import AimCallback, AimTracker  # noqa F401

aim/keras_tuner.py

@@ -0,0 +1,2 @@
+# Alias to SDK Keras-Tuner interface
+from aim.sdk.adapters.keras_tuner import AimCallback  # noqa F401

aim/lightgbm.py

@@ -0,0 +1,2 @@
+# Alias to SDK LightGBM interface
+from aim.sdk.adapters.lightgbm import AimCallback  # noqa: F401

aim/mxnet.py

@@ -0,0 +1,2 @@
+# Alias to SDK mxnet interface
+from aim.sdk.adapters.mxnet import AimLoggingHandler  # noqa F401

aim/optuna.py

@@ -0,0 +1,2 @@
+# Alias to SDK Optuna interface
+from aim.sdk.adapters.optuna import AimCallback  # noqa F401

aim/paddle.py

@@ -0,0 +1,2 @@
+# Alias to SDK PaddlePaddle interface
+from aim.sdk.adapters.paddle import AimCallback  # noqa F401

aim/prophet.py

@@ -0,0 +1,2 @@
+# Alias to SDK Prophet interface
+from aim.sdk.adapters.prophet import AimLogger  # noqa F401

aim/pytorch.py

@@ -0,0 +1,2 @@
+# Alias to SDK PyTorch utils
+from aim.sdk.adapters.pytorch import track_params_dists, track_gradients_dists # noqa