Create migrate-users function

Creating a new cometx/cli/migrate_users.py module, which takes a chargeback report from a source environment, and invites users to all of the same workspaces in a destination environment

Adding --create-workspaces argument to both copy and migreate_users so user can decide whether to create a workspace when it doesn't exist.

Author: chasefortier

README.md

@@ -46,6 +46,7 @@ for ML management.
 * [cometx download](#cometx-download)
 * [cometx list](#cometx-list)
 * [cometx log](#cometx-log)
+* [cometx migrate-users](#cometx-migrate-users)
 * [cometx rename-duplicates](#cometx-rename-duplicates)
 * [cometx reproduce](#cometx-reproduce)
 * [cometx smoke-test](#cometx-smoke-test)
@@ -268,6 +269,7 @@ Not all combinations are possible:
 * `--symlink` - Instead of copying, create a link to an experiment in a project
 * `--sync` - Check to see if experiment name has been created first; if so, skip
 * `--path PATH` - Path to prepend to workspace_src when accessing files (supports ~ for home directory)
+* `--create-workspaces` - Attempt to create the destination workspace if it does not exist (requires appropriate API permissions)
 
 ### Using --path
 
@@ -329,6 +331,43 @@ Where TYPE is one of the following names:
 
 For more information, `cometx log --help`
 
+## cometx migrate-users
+
+This command is used to migrate users into workspaces from a source Comet environment to a destination environment. It reads workspace membership from a chargeback report from the source environment and invites each user to the corresponding workspace in the destination environment by email. If a user with the email does not exist in the destination environment, they still will be provisioned access to these workspaces once they sign up using that email.
+
+The chargeback report can be fetched automatically from the source environment's admin API, or you can provide a pre-downloaded JSON file.
+
+```
+cometx migrate-users --api-key DEST_KEY --source-api-key SOURCE_KEY [FLAGS ...]
+cometx migrate-users --api-key DEST_KEY --chargeback-report /path/to/report.json [FLAGS ...]
+```
+
+**Arguments:**
+* `--api-key API_KEY` - API key for the destination environment where users will be added. Falls back to `COMET_API_KEY` environment variable if not provided.
+* `--source-api-key SOURCE_API_KEY` - API key for the source environment. Used to fetch the chargeback report. Falls back to `COMET_API_KEY` environment variable if not provided. Required unless `--chargeback-report` is given.
+* `--chargeback-report PATH` - Path to a local chargeback report JSON file. When provided, the report is loaded from this file instead of being fetched from the source environment, and `--source-api-key` is not required.
+
+### Flags
+
+* `--create-workspaces` - Create workspaces on the destination environment if they don't already exist (default: off)
+* `--dry-run` - Print what would happen without making any changes
+
+### Examples
+
+```bash
+# Dry run — preview what would happen
+cometx migrate-users --api-key DEST_KEY --source-api-key SOURCE_KEY --dry-run
+
+# Execute the migration
+cometx migrate-users --api-key DEST_KEY --source-api-key SOURCE_KEY
+
+# Use a local chargeback report file
+cometx migrate-users --api-key DEST_KEY --chargeback-report /tmp/chargeback_reports.json
+
+```
+
+For more information, `cometx migrate-users --help`
+
 ## cometx rename-duplicates
 
 This command is used to rename duplicate experiments within projects. When multiple experiments share the same name in a project, this command renames the duplicates to NAME-1, NAME-2, etc. while avoiding conflicts with existing names.

cometx/cli/__init__.py

@@ -26,6 +26,7 @@
     cometx smoke-test
     cometx update
     cometx admin
+    cometx migrate-users
 
 For more information:
     cometx COMMAND --help
@@ -47,6 +48,7 @@
     download,
     list_command,
     log,
+    migrate_users,
     rename_duplicates,
     reproduce,
     smoke_test,
@@ -108,6 +110,7 @@ def main(raw_args=sys.argv[1:]):
     add_subparser(subparsers, config, "config")
     add_subparser(subparsers, rename_duplicates, "rename-duplicates")
     add_subparser(subparsers, smoke_test, "smoke-test")
+    add_subparser(subparsers, migrate_users, "migrate-users")
     add_subparser(
         subparsers, upload_optimizer_experiments, "upload-optimizer-experiments"
     )

cometx/cli/copy.py

@@ -197,6 +197,12 @@ def get_parser_arguments(parser):
         type=str,
         default=None,
     )
+    parser.add_argument(
+        "--create-workspaces",
+        help="Attempt to create the destination workspace if it does not exist",
+        default=False,
+        action="store_true",
+    )
 
 
 def copy(parsed_args, remaining=None):
@@ -214,6 +220,7 @@ def copy(parsed_args, remaining=None):
             parsed_args.ignore,
             parsed_args.debug,
             parsed_args.sync,
+            parsed_args.create_workspaces,
         )
 
         # Wait for all uploads to complete
@@ -790,7 +797,7 @@ def cleanup(self):
         # Restore default signal handler
         signal.signal(signal.SIGINT, signal.default_int_handler)
 
-    def copy(self, source, destination, symlink, ignore, debug, sync):
+    def copy(self, source, destination, symlink, ignore, debug, sync, create_workspaces=False):
         """ """
         self.ignore = ignore
         self.debug = debug
@@ -820,31 +827,37 @@ def copy(self, source, destination, symlink, ignore, debug, sync):
         else:
             raise Exception("invalid COMET_SOURCE: %r" % source)
 
-        # First check to make sure workspace_dst exists; try to create if not:
+        # First check to make sure workspace_dst exists:
         workspaces = self.api.get_workspaces()
         if workspace_dst not in workspaces:
-            print(
-                f"Workspace {workspace_dst!r} does not exist, attempting to create it..."
-            )
-            try:
-                import requests
-
-                url = f"{self.api.server_url}/api/rest/v2/write/workspace/new"
-                response = requests.post(
-                    url,
-                    json={"name": workspace_dst},
-                    headers={
-                        "Authorization": self.api.api_key,
-                        "Content-Type": "application/json",
-                    },
+            if create_workspaces:
+                print(
+                    f"Workspace {workspace_dst!r} does not exist, attempting to create it..."
                 )
-                response.raise_for_status()
-                print(f"Workspace {workspace_dst!r} created successfully.")
-            except Exception as exc:
+                try:
+                    import requests
+
+                    url = f"{self.api.server_url}/api/rest/v2/write/workspace/new"
+                    response = requests.post(
+                        url,
+                        json={"name": workspace_dst},
+                        headers={
+                            "Authorization": self.api.api_key,
+                            "Content-Type": "application/json",
+                        },
+                    )
+                    response.raise_for_status()
+                    print(f"Workspace {workspace_dst!r} created successfully.")
+                except Exception as exc:
+                    raise Exception(
+                        f"Workspace {workspace_dst!r} does not exist and could not be "
+                        f"created automatically: {exc}. "
+                        f"Please create it via the Comet UI and try again."
+                    )
+            else:
                 raise Exception(
-                    f"Workspace {workspace_dst!r} does not exist and could not be "
-                    f"created automatically: {exc}. "
-                    f"Please create it via the Comet UI or with an admin API Key and try again."
+                    f"{workspace_dst} does not exist; use --create-workspaces to "
+                    f"create it automatically, or create it via the Comet UI"
                 )
 
         if project_src == "panels":