Add workspace creation option to copy + new migrate-users command

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,47 @@ Where TYPE is one of the following names:
 
 For more information, `cometx log --help`
 
+## cometx migrate-users
+
+This command migrates users into workspaces from a source Comet environment to a destination environment. It reads workspace membership from a chargeback report fetched from the source environment's admin API (or a pre-downloaded local file) and invites each user to the corresponding workspace in the destination environment by email. If a user with that email does not yet exist in the destination, they will be provisioned access once they sign up.
+
+```
+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. Falls back to `COMET_API_KEY` if not provided.
+* `--url URL` - Base URL of the destination Comet environment (e.g. `https://comet.example.com`). Required for self-hosted instances when the API key does not encode the server URL.
+* `--source-api-key SOURCE_API_KEY` - API key for the source environment. Used to fetch the chargeback report. Required unless `--chargeback-report` is given.
+* `--source-url SOURCE_URL` - Base URL of the source Comet environment. Required for self-hosted source instances when the source API key does not encode the server URL.
+* `--chargeback-report PATH` - Path to a local chargeback report JSON file. When provided, `--source-api-key` is not required.
+
+### Flags
+
+* `--create-workspaces` - Create workspaces on the destination if they don't already exist (default: off)
+* `--dry-run` - Print what would happen without making any changes
+* `--failures-output PATH` - Path to write failed operations JSON (default: `bulk_add_failures_by_email.json`)
+
+### 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
+
+# Self-hosted environments — provide explicit URLs
+cometx migrate-users --api-key DEST_KEY --url https://comet.dest.example.com \
+    --source-api-key SOURCE_KEY --source-url https://comet.src.example.com
+```
+
+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.
@@ -479,12 +522,8 @@ cometx admin chargeback-report [YEAR-MONTH]
 **Examples:**
 ```
 cometx admin chargeback-report
-<<<<<<< Updated upstream
-cometx admin usage-report
-=======
 cometx admin chargeback-report 2024-09
 ```
->>>>>>> Stashed changes
 
 #### usage-report
 

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

@@ -78,6 +78,8 @@
 import time
 import urllib.parse
 import zipfile
+
+import requests
 from datetime import datetime, timedelta
 
 from comet_ml import APIExperiment, Artifact, Experiment, OfflineExperiment
@@ -102,6 +104,7 @@
 from ..framework.comet.download_manager import sanitize_filename
 from ..utils import remove_extra_slashes
 from .copy_utils import upload_single_offline_experiment
+from .migrate_users import _create_workspace
 
 ADDITIONAL_ARGS = False
 
@@ -197,6 +200,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 +223,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 +800,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
@@ -823,9 +833,35 @@ def copy(self, source, destination, symlink, ignore, debug, sync):
         # First check to make sure workspace_dst exists:
         workspaces = self.api.get_workspaces()
         if workspace_dst not in workspaces:
-            raise Exception(
-                f"{workspace_dst} does not exist; use the Comet UI to create it"
-            )
+            if create_workspaces:
+                print(
+                    f"Workspace {workspace_dst!r} does not exist, attempting to create it..."
+                )
+                dest_url = self.api._get_url_server()
+                headers = {
+                    "Authorization": self.api.api_key,
+                    "Content-Type": "application/json",
+                }
+                try:
+                    _create_workspace(dest_url, headers, workspace_dst)
+                    print(f"Workspace {workspace_dst!r} created successfully.")
+                except requests.exceptions.HTTPError as exc:
+                    raise Exception(
+                        f"Workspace {workspace_dst!r} does not exist and could not be "
+                        f"created automatically (HTTP {exc.response.status_code}). "
+                        f"Please create it via the Comet UI and try again."
+                    ) from exc
+                except requests.exceptions.RequestException 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."
+                    ) from exc
+            else:
+                raise Exception(
+                    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":
             # experiment_src may be "*" or filename